Actual source code: veckok.kokkos.cxx

      1: /*
      2:    Implements the sequential Kokkos vectors.
      3: */
      4: #include <petsc_kokkos.hpp>
      5: #include <petscvec_kokkos.hpp>

      7: #include <petsc/private/sfimpl.h>
      8: #include <petsc/private/petscimpl.h>
      9: #include <petscmath.h>
     10: #include <petscviewer.h>
     11: #include <KokkosBlas.hpp>
     12: #include <Kokkos_Functional.hpp>

     14: #include <../src/vec/vec/impls/dvecimpl.h>
     15: #include <../src/vec/vec/impls/seq/kokkos/veckokkosimpl.hpp>

     17: template <class MemorySpace>
     18: static PetscErrorCode VecGetKokkosView_Private(Vec v, PetscScalarKokkosViewType<MemorySpace> *kv, PetscBool overwrite)
     19: {
     20:   Vec_Kokkos *veckok   = static_cast<Vec_Kokkos *>(v->spptr);
     21:   using ExecutionSpace = typename PetscScalarKokkosViewType<MemorySpace>::traits::device_type;

     23:   PetscFunctionBegin;
     24:   VecErrorIfNotKokkos(v);
     25:   if (!overwrite) { /* If overwrite=true, no need to sync the space, since caller will overwrite the data */
     26:     PetscCallCXX(veckok->v_dual.sync<ExecutionSpace>());
     27:   }
     28:   PetscCallCXX(*kv = veckok->v_dual.view<ExecutionSpace>());
     29:   PetscFunctionReturn(PETSC_SUCCESS);
     30: }

     32: template <class MemorySpace>
     33: static PetscErrorCode VecRestoreKokkosView_Private(Vec v, PetscScalarKokkosViewType<MemorySpace> *, PetscBool overwrite)
     34: {
     35:   Vec_Kokkos *veckok   = static_cast<Vec_Kokkos *>(v->spptr);
     36:   using ExecutionSpace = typename PetscScalarKokkosViewType<MemorySpace>::traits::device_type;

     38:   PetscFunctionBegin;
     39:   VecErrorIfNotKokkos(v);
     40:   if (overwrite) PetscCallCXX(veckok->v_dual.clear_sync_state()); /* If overwrite=true, clear the old sync state since user forced an overwrite */
     41:   PetscCallCXX(veckok->v_dual.modify<ExecutionSpace>());
     42:   PetscCall(PetscObjectStateIncrease((PetscObject)v));
     43:   PetscFunctionReturn(PETSC_SUCCESS);
     44: }

     46: template <class MemorySpace>
     47: PetscErrorCode VecGetKokkosView(Vec v, ConstPetscScalarKokkosViewType<MemorySpace> *kv)
     48: {
     49:   Vec_Kokkos *veckok   = static_cast<Vec_Kokkos *>(v->spptr);
     50:   using ExecutionSpace = typename PetscScalarKokkosViewType<MemorySpace>::traits::device_type;

     52:   PetscFunctionBegin;
     53:   VecErrorIfNotKokkos(v);
     54:   PetscCallCXX(veckok->v_dual.sync<ExecutionSpace>());
     55:   PetscCallCXX(*kv = veckok->v_dual.view<ExecutionSpace>());
     56:   PetscFunctionReturn(PETSC_SUCCESS);
     57: }

     59: /* Function template explicit instantiation */
     60: template PETSC_VISIBILITY_PUBLIC PetscErrorCode VecGetKokkosView(Vec, ConstPetscScalarKokkosView *);
     61: template <>
     62: PETSC_VISIBILITY_PUBLIC PetscErrorCode VecGetKokkosView(Vec v, PetscScalarKokkosView *kv)
     63: {
     64:   return VecGetKokkosView_Private(v, kv, PETSC_FALSE);
     65: }
     66: template <>
     67: PETSC_VISIBILITY_PUBLIC PetscErrorCode VecRestoreKokkosView(Vec v, PetscScalarKokkosView *kv)
     68: {
     69:   return VecRestoreKokkosView_Private(v, kv, PETSC_FALSE);
     70: }
     71: template <>
     72: PETSC_VISIBILITY_PUBLIC PetscErrorCode VecGetKokkosViewWrite(Vec v, PetscScalarKokkosView *kv)
     73: {
     74:   return VecGetKokkosView_Private(v, kv, PETSC_TRUE);
     75: }
     76: template <>
     77: PETSC_VISIBILITY_PUBLIC PetscErrorCode VecRestoreKokkosViewWrite(Vec v, PetscScalarKokkosView *kv)
     78: {
     79:   return VecRestoreKokkosView_Private(v, kv, PETSC_TRUE);
     80: }

     82: #if !defined(KOKKOS_ENABLE_UNIFIED_MEMORY) /* Get host views if the default memory space is not host space */
     83: template PETSC_VISIBILITY_PUBLIC PetscErrorCode VecGetKokkosView(Vec, ConstPetscScalarKokkosViewHost *);
     84: template <>
     85: PETSC_VISIBILITY_PUBLIC PetscErrorCode VecGetKokkosView(Vec v, PetscScalarKokkosViewHost *kv)
     86: {
     87:   return VecGetKokkosView_Private(v, kv, PETSC_FALSE);
     88: }
     89: template <>
     90: PETSC_VISIBILITY_PUBLIC PetscErrorCode VecRestoreKokkosView(Vec v, PetscScalarKokkosViewHost *kv)
     91: {
     92:   return VecRestoreKokkosView_Private(v, kv, PETSC_FALSE);
     93: }
     94: template <>
     95: PETSC_VISIBILITY_PUBLIC PetscErrorCode VecGetKokkosViewWrite(Vec v, PetscScalarKokkosViewHost *kv)
     96: {
     97:   return VecGetKokkosView_Private(v, kv, PETSC_TRUE);
     98: }
     99: template <>
    100: PETSC_VISIBILITY_PUBLIC PetscErrorCode VecRestoreKokkosViewWrite(Vec v, PetscScalarKokkosViewHost *kv)
    101: {
    102:   return VecRestoreKokkosView_Private(v, kv, PETSC_TRUE);
    103: }
    104: #endif

    106: PetscErrorCode VecSetRandom_SeqKokkos(Vec xin, PetscRandom r)
    107: {
    108:   const PetscInt n = xin->map->n;
    109:   PetscScalar   *xx;

    111:   PetscFunctionBegin;
    112:   PetscCall(VecGetArrayWrite(xin, &xx)); /* TODO: generate randoms directly on device */
    113:   for (PetscInt i = 0; i < n; i++) PetscCall(PetscRandomGetValue(r, &xx[i]));
    114:   PetscCall(VecRestoreArrayWrite(xin, &xx));
    115:   PetscFunctionReturn(PETSC_SUCCESS);
    116: }

    118: /* x = |x| */
    119: PetscErrorCode VecAbs_SeqKokkos(Vec xin)
    120: {
    121:   PetscScalarKokkosView xv;
    122:   auto                  exec = PetscGetKokkosExecutionSpace();

    124:   PetscFunctionBegin;
    125:   PetscCall(PetscLogGpuTimeBegin());
    126:   PetscCall(VecGetKokkosView(xin, &xv));
    127:   PetscCallCXX(KokkosBlas::abs(exec, xv, xv));
    128:   PetscCall(VecRestoreKokkosView(xin, &xv));
    129:   PetscCall(PetscLogGpuTimeEnd());
    130:   PetscFunctionReturn(PETSC_SUCCESS);
    131: }

    133: /* x = 1/x */
    134: PetscErrorCode VecReciprocal_SeqKokkos(Vec xin)
    135: {
    136:   PetscScalarKokkosView xv;

    138:   PetscFunctionBegin;
    139:   PetscCall(PetscLogGpuTimeBegin());
    140:   PetscCall(VecGetKokkosView(xin, &xv));
    141:   PetscCallCXX(Kokkos::parallel_for(
    142:     "VecReciprocal", Kokkos::RangePolicy<>(PetscGetKokkosExecutionSpace(), 0, xin->map->n), KOKKOS_LAMBDA(const PetscInt &i) {
    143:       if (xv(i) != (PetscScalar)0.0) xv(i) = (PetscScalar)1.0 / xv(i);
    144:     }));
    145:   PetscCall(VecRestoreKokkosView(xin, &xv));
    146:   PetscCall(PetscLogGpuTimeEnd());
    147:   PetscFunctionReturn(PETSC_SUCCESS);
    148: }

    150: PetscErrorCode VecMin_SeqKokkos(Vec xin, PetscInt *p, PetscReal *val)
    151: {
    152:   ConstPetscScalarKokkosView                           xv;
    153:   Kokkos::MinFirstLoc<PetscReal, PetscInt>::value_type result;

    155:   PetscFunctionBegin;
    156:   PetscCall(PetscLogGpuTimeBegin());
    157:   PetscCall(VecGetKokkosView(xin, &xv));
    158:   PetscCallCXX(Kokkos::parallel_reduce(
    159:     "VecMin", Kokkos::RangePolicy<>(PetscGetKokkosExecutionSpace(), 0, xin->map->n),
    160:     KOKKOS_LAMBDA(const PetscInt &i, Kokkos::MinFirstLoc<PetscReal, PetscInt>::value_type &lupdate) {
    161:       if (PetscRealPart(xv(i)) < lupdate.val) {
    162:         lupdate.val = PetscRealPart(xv(i));
    163:         lupdate.loc = i;
    164:       }
    165:     },
    166:     Kokkos::MinFirstLoc<PetscReal, PetscInt>(result)));
    167:   *val = result.val;
    168:   if (p) *p = result.loc;
    169:   PetscCall(VecRestoreKokkosView(xin, &xv));
    170:   PetscCall(PetscLogGpuTimeEnd());
    171:   PetscFunctionReturn(PETSC_SUCCESS);
    172: }

    174: PetscErrorCode VecMax_SeqKokkos(Vec xin, PetscInt *p, PetscReal *val)
    175: {
    176:   ConstPetscScalarKokkosView                           xv;
    177:   Kokkos::MaxFirstLoc<PetscReal, PetscInt>::value_type result;

    179:   PetscFunctionBegin;
    180:   PetscCall(PetscLogGpuTimeBegin());
    181:   PetscCall(VecGetKokkosView(xin, &xv));
    182:   PetscCallCXX(Kokkos::parallel_reduce(
    183:     "VecMax", Kokkos::RangePolicy<>(PetscGetKokkosExecutionSpace(), 0, xin->map->n),
    184:     KOKKOS_LAMBDA(const PetscInt &i, Kokkos::MaxFirstLoc<PetscReal, PetscInt>::value_type &lupdate) {
    185:       if (PetscRealPart(xv(i)) > lupdate.val) {
    186:         lupdate.val = PetscRealPart(xv(i));
    187:         lupdate.loc = i;
    188:       }
    189:     },
    190:     Kokkos::MaxFirstLoc<PetscReal, PetscInt>(result)));
    191:   *val = result.val;
    192:   if (p) *p = result.loc;
    193:   PetscCall(VecRestoreKokkosView(xin, &xv));
    194:   PetscCall(PetscLogGpuTimeEnd());
    195:   PetscFunctionReturn(PETSC_SUCCESS);
    196: }

    198: PetscErrorCode VecSum_SeqKokkos(Vec xin, PetscScalar *sum)
    199: {
    200:   ConstPetscScalarKokkosView xv;

    202:   PetscFunctionBegin;
    203:   PetscCall(PetscLogGpuTimeBegin());
    204:   PetscCall(VecGetKokkosView(xin, &xv));
    205:   PetscCallCXX(*sum = KokkosBlas::sum(PetscGetKokkosExecutionSpace(), xv));
    206:   PetscCall(VecRestoreKokkosView(xin, &xv));
    207:   PetscCall(PetscLogGpuTimeEnd());
    208:   PetscFunctionReturn(PETSC_SUCCESS);
    209: }

    211: PetscErrorCode VecShift_SeqKokkos(Vec xin, PetscScalar shift)
    212: {
    213:   PetscScalarKokkosView xv;

    215:   PetscFunctionBegin;
    216:   PetscCall(PetscLogGpuTimeBegin());
    217:   PetscCall(VecGetKokkosView(xin, &xv));
    218:   PetscCallCXX(Kokkos::parallel_for("VecShift", Kokkos::RangePolicy<>(PetscGetKokkosExecutionSpace(), 0, xin->map->n), KOKKOS_LAMBDA(const PetscInt &i) { xv(i) += shift; }); PetscCall(VecRestoreKokkosView(xin, &xv)));
    219:   PetscCall(PetscLogGpuTimeEnd());
    220:   PetscCall(PetscLogGpuFlops(xin->map->n));
    221:   PetscFunctionReturn(PETSC_SUCCESS);
    222: }

    224: /* y = alpha x + y */
    225: PetscErrorCode VecAXPY_SeqKokkos(Vec yin, PetscScalar alpha, Vec xin)
    226: {
    227:   PetscFunctionBegin;
    228:   if (alpha == (PetscScalar)0.0) PetscFunctionReturn(PETSC_SUCCESS);
    229:   if (yin == xin) PetscCall(VecScale_SeqKokkos(yin, alpha + 1));
    230:   else {
    231:     PetscBool xiskok, yiskok;

    233:     PetscCall(PetscObjectTypeCompareAny((PetscObject)xin, &xiskok, VECSEQKOKKOS, VECMPIKOKKOS, ""));
    234:     PetscCall(PetscObjectTypeCompareAny((PetscObject)yin, &yiskok, VECSEQKOKKOS, VECMPIKOKKOS, ""));
    235:     if (xiskok && yiskok) {
    236:       PetscScalarKokkosView      yv;
    237:       ConstPetscScalarKokkosView xv;

    239:       PetscCall(PetscLogGpuTimeBegin());
    240:       PetscCall(VecGetKokkosView(xin, &xv));
    241:       PetscCall(VecGetKokkosView(yin, &yv));
    242:       PetscCallCXX(KokkosBlas::axpy(PetscGetKokkosExecutionSpace(), alpha, xv, yv));
    243:       PetscCall(VecRestoreKokkosView(xin, &xv));
    244:       PetscCall(VecRestoreKokkosView(yin, &yv));
    245:       PetscCall(PetscLogGpuTimeEnd());
    246:       PetscCall(PetscLogGpuFlops(2.0 * yin->map->n));
    247:     } else {
    248:       PetscCall(VecAXPY_Seq(yin, alpha, xin));
    249:     }
    250:   }
    251:   PetscFunctionReturn(PETSC_SUCCESS);
    252: }

    254: /* y = x + beta y */
    255: PetscErrorCode VecAYPX_SeqKokkos(Vec yin, PetscScalar beta, Vec xin)
    256: {
    257:   PetscFunctionBegin;
    258:   /* One needs to define KOKKOSBLAS_OPTIMIZATION_LEVEL_AXPBY > 2 to have optimizations for cases alpha/beta = 0,+/-1 */
    259:   PetscCall(VecAXPBY_SeqKokkos(yin, 1.0, beta, xin));
    260:   PetscFunctionReturn(PETSC_SUCCESS);
    261: }

    263: /* z = y^T x */
    264: PetscErrorCode VecTDot_SeqKokkos(Vec xin, Vec yin, PetscScalar *z)
    265: {
    266:   ConstPetscScalarKokkosView xv, yv;

    268:   PetscFunctionBegin;
    269:   PetscCall(PetscLogGpuTimeBegin());
    270:   PetscCall(VecGetKokkosView(xin, &xv));
    271:   PetscCall(VecGetKokkosView(yin, &yv));
    272:   // Kokkos always overwrites z, so no need to init it
    273:   PetscCallCXX(Kokkos::parallel_reduce("VecTDot", Kokkos::RangePolicy<>(PetscGetKokkosExecutionSpace(), 0, xin->map->n), KOKKOS_LAMBDA(const PetscInt &i, PetscScalar &update) { update += yv(i) * xv(i); }, *z));
    274:   PetscCall(VecRestoreKokkosView(yin, &yv));
    275:   PetscCall(VecRestoreKokkosView(xin, &xv));
    276:   PetscCall(PetscLogGpuTimeEnd());
    277:   if (xin->map->n > 0) PetscCall(PetscLogGpuFlops(2.0 * xin->map->n));
    278:   PetscFunctionReturn(PETSC_SUCCESS);
    279: }

    281: struct TransposeDotTag { };
    282: struct ConjugateDotTag { };

    284: template <PetscInt ValueCount>
    285: struct MDotFunctor {
    286:   static_assert(ValueCount >= 1 && ValueCount <= 8, "ValueCount must be in [1, 8]");
    287:   /* Note the C++ notation for an array typedef */
    288:   // noted, thanks
    289:   typedef PetscScalar                           value_type[];
    290:   typedef ConstPetscScalarKokkosView::size_type size_type;

    292:   /* Tell Kokkos the result array's number of entries. This must be a public value in the functor */
    293:   static constexpr size_type value_count = ValueCount;
    294:   ConstPetscScalarKokkosView xv, yv[8];

    296:   MDotFunctor(ConstPetscScalarKokkosView &xv, ConstPetscScalarKokkosView &yv0, ConstPetscScalarKokkosView &yv1, ConstPetscScalarKokkosView &yv2, ConstPetscScalarKokkosView &yv3, ConstPetscScalarKokkosView &yv4, ConstPetscScalarKokkosView &yv5, ConstPetscScalarKokkosView &yv6, ConstPetscScalarKokkosView &yv7) :
    297:     xv(xv)
    298:   {
    299:     yv[0] = yv0;
    300:     yv[1] = yv1;
    301:     yv[2] = yv2;
    302:     yv[3] = yv3;
    303:     yv[4] = yv4;
    304:     yv[5] = yv5;
    305:     yv[6] = yv6;
    306:     yv[7] = yv7;
    307:   }

    309:   KOKKOS_INLINE_FUNCTION void operator()(TransposeDotTag, const size_type i, value_type sum) const
    310:   {
    311:     PetscScalar xval = xv(i);
    312:     for (size_type j = 0; j < value_count; ++j) sum[j] += yv[j](i) * xval;
    313:   }

    315:   KOKKOS_INLINE_FUNCTION void operator()(ConjugateDotTag, const size_type i, value_type sum) const
    316:   {
    317:     PetscScalar xval = xv(i);
    318:     for (size_type j = 0; j < value_count; ++j) sum[j] += PetscConj(yv[j](i)) * xval;
    319:   }

    321:   // Per https://kokkos.github.io/kokkos-core-wiki/API/core/parallel-dispatch/parallel_reduce.html#requirements
    322:   // "when specifying a tag in the policy, the functor's potential init/join/final member functions must also be tagged"
    323:   // So we have this kind of duplicated code.
    324:   KOKKOS_INLINE_FUNCTION void join(TransposeDotTag, value_type dst, const value_type src) const { join(dst, src); }
    325:   KOKKOS_INLINE_FUNCTION void join(ConjugateDotTag, value_type dst, const value_type src) const { join(dst, src); }

    327:   KOKKOS_INLINE_FUNCTION void init(TransposeDotTag, value_type sum) const { init(sum); }
    328:   KOKKOS_INLINE_FUNCTION void init(ConjugateDotTag, value_type sum) const { init(sum); }

    330:   KOKKOS_INLINE_FUNCTION void join(value_type dst, const value_type src) const
    331:   {
    332:     for (size_type j = 0; j < value_count; j++) dst[j] += src[j];
    333:   }

    335:   KOKKOS_INLINE_FUNCTION void init(value_type sum) const
    336:   {
    337:     for (size_type j = 0; j < value_count; j++) sum[j] = 0.0;
    338:   }
    339: };

    341: template <class WorkTag>
    342: PetscErrorCode VecMultiDot_Private(Vec xin, PetscInt nv, const Vec yin[], PetscScalar *z)
    343: {
    344:   PetscInt                   i, j, cur = 0, ngroup = nv / 8, rem = nv % 8, N = xin->map->n;
    345:   ConstPetscScalarKokkosView xv, yv[8];
    346:   PetscScalarKokkosViewHost  zv(z, nv);
    347:   auto                       exec = PetscGetKokkosExecutionSpace();

    349:   PetscFunctionBegin;
    350:   PetscCall(VecGetKokkosView(xin, &xv));
    351:   for (i = 0; i < ngroup; i++) { /* 8 y's per group */
    352:     for (j = 0; j < 8; j++) PetscCall(VecGetKokkosView(yin[cur + j], &yv[j]));
    353:     MDotFunctor<8> mdot(xv, yv[0], yv[1], yv[2], yv[3], yv[4], yv[5], yv[6], yv[7]); /* Hope Kokkos make it asynchronous */
    354:     PetscCall(PetscLogGpuTimeBegin());
    355:     PetscCallCXX(Kokkos::parallel_reduce(Kokkos::RangePolicy<WorkTag>(exec, 0, N), mdot, Kokkos::subview(zv, Kokkos::pair<PetscInt, PetscInt>(cur, cur + 8))));
    356:     PetscCall(PetscLogGpuTimeEnd());
    357:     for (j = 0; j < 8; j++) PetscCall(VecRestoreKokkosView(yin[cur + j], &yv[j]));
    358:     cur += 8;
    359:   }

    361:   if (rem) { /* The remaining */
    362:     for (j = 0; j < rem; j++) PetscCall(VecGetKokkosView(yin[cur + j], &yv[j]));
    363:     Kokkos::RangePolicy<WorkTag> policy(exec, 0, N);
    364:     auto                         results = Kokkos::subview(zv, Kokkos::pair<PetscInt, PetscInt>(cur, cur + rem));
    365:     // clang-format off
    366:     PetscCall(PetscLogGpuTimeBegin());
    367:     switch (rem) {
    368:     case 1: PetscCallCXX(Kokkos::parallel_reduce(policy, MDotFunctor<1>(xv, yv[0], yv[1], yv[2], yv[3], yv[4], yv[5], yv[6], yv[7]), results)); break;
    369:     case 2: PetscCallCXX(Kokkos::parallel_reduce(policy, MDotFunctor<2>(xv, yv[0], yv[1], yv[2], yv[3], yv[4], yv[5], yv[6], yv[7]), results)); break;
    370:     case 3: PetscCallCXX(Kokkos::parallel_reduce(policy, MDotFunctor<3>(xv, yv[0], yv[1], yv[2], yv[3], yv[4], yv[5], yv[6], yv[7]), results)); break;
    371:     case 4: PetscCallCXX(Kokkos::parallel_reduce(policy, MDotFunctor<4>(xv, yv[0], yv[1], yv[2], yv[3], yv[4], yv[5], yv[6], yv[7]), results)); break;
    372:     case 5: PetscCallCXX(Kokkos::parallel_reduce(policy, MDotFunctor<5>(xv, yv[0], yv[1], yv[2], yv[3], yv[4], yv[5], yv[6], yv[7]), results)); break;
    373:     case 6: PetscCallCXX(Kokkos::parallel_reduce(policy, MDotFunctor<6>(xv, yv[0], yv[1], yv[2], yv[3], yv[4], yv[5], yv[6], yv[7]), results)); break;
    374:     case 7: PetscCallCXX(Kokkos::parallel_reduce(policy, MDotFunctor<7>(xv, yv[0], yv[1], yv[2], yv[3], yv[4], yv[5], yv[6], yv[7]), results)); break;
    375:     }
    376:     PetscCall(PetscLogGpuTimeEnd());
    377:     // clang-format on
    378:     for (j = 0; j < rem; j++) PetscCall(VecRestoreKokkosView(yin[cur + j], &yv[j]));
    379:   }
    380:   PetscCall(VecRestoreKokkosView(xin, &xv));
    381:   exec.fence(); /* If reduce is async, then we need this fence to make sure z is ready for use on host */
    382:   PetscFunctionReturn(PETSC_SUCCESS);
    383: }

    385: static PetscErrorCode VecMultiDot_Verbose(Vec xin, PetscInt nv, const Vec yin[], PetscScalar *z)
    386: {
    387:   PetscInt                   ngroup = nv / 8, rem = nv % 8, N = xin->map->n;
    388:   ConstPetscScalarKokkosView xv, y0, y1, y2, y3, y4, y5, y6, y7;
    389:   PetscScalar               *zp = z;
    390:   const Vec                 *yp = yin;
    391:   Kokkos::RangePolicy<>      policy(PetscGetKokkosExecutionSpace(), 0, N);

    393:   // clang-format off
    394:   PetscFunctionBegin;
    395:   PetscCall(VecGetKokkosView(xin, &xv));
    396:   for (PetscInt k = 0; k < ngroup; k++) { // 8 y's per group
    397:     PetscCall(VecGetKokkosView(yp[0], &y0));
    398:     PetscCall(VecGetKokkosView(yp[1], &y1));
    399:     PetscCall(VecGetKokkosView(yp[2], &y2));
    400:     PetscCall(VecGetKokkosView(yp[3], &y3));
    401:     PetscCall(VecGetKokkosView(yp[4], &y4));
    402:     PetscCall(VecGetKokkosView(yp[5], &y5));
    403:     PetscCall(VecGetKokkosView(yp[6], &y6));
    404:     PetscCall(VecGetKokkosView(yp[7], &y7));
    405:     PetscCall(PetscLogGpuTimeBegin()); // only for GPU kernel execution
    406:     Kokkos::parallel_reduce(
    407:       "VecMDot8", policy,
    408:       KOKKOS_LAMBDA(const PetscInt &i, PetscScalar &lsum0, PetscScalar &lsum1, PetscScalar &lsum2, PetscScalar &lsum3, PetscScalar &lsum4, PetscScalar &lsum5, PetscScalar &lsum6, PetscScalar &lsum7) {
    409:         lsum0 += xv(i) * PetscConj(y0(i)); lsum1 += xv(i) * PetscConj(y1(i)); lsum2 += xv(i) * PetscConj(y2(i)); lsum3 += xv(i) * PetscConj(y3(i));
    410:         lsum4 += xv(i) * PetscConj(y4(i)); lsum5 += xv(i) * PetscConj(y5(i)); lsum6 += xv(i) * PetscConj(y6(i)); lsum7 += xv(i) * PetscConj(y7(i));
    411:       }, zp[0], zp[1], zp[2], zp[3], zp[4], zp[5], zp[6], zp[7]);
    412:     PetscCall(PetscLogGpuTimeEnd());
    413:     PetscCall(PetscLogGpuToCpu(8 * sizeof(PetscScalar))); // for copying to z[] on host
    414:     PetscCall(VecRestoreKokkosView(yp[0], &y0));
    415:     PetscCall(VecRestoreKokkosView(yp[1], &y1));
    416:     PetscCall(VecRestoreKokkosView(yp[2], &y2));
    417:     PetscCall(VecRestoreKokkosView(yp[3], &y3));
    418:     PetscCall(VecRestoreKokkosView(yp[4], &y4));
    419:     PetscCall(VecRestoreKokkosView(yp[5], &y5));
    420:     PetscCall(VecRestoreKokkosView(yp[6], &y6));
    421:     PetscCall(VecRestoreKokkosView(yp[7], &y7));
    422:     yp += 8;
    423:     zp += 8;
    424:   }

    426:   if (rem) { /* The remaining */
    427:     if (rem > 0) PetscCall(VecGetKokkosView(yp[0], &y0));
    428:     if (rem > 1) PetscCall(VecGetKokkosView(yp[1], &y1));
    429:     if (rem > 2) PetscCall(VecGetKokkosView(yp[2], &y2));
    430:     if (rem > 3) PetscCall(VecGetKokkosView(yp[3], &y3));
    431:     if (rem > 4) PetscCall(VecGetKokkosView(yp[4], &y4));
    432:     if (rem > 5) PetscCall(VecGetKokkosView(yp[5], &y5));
    433:     if (rem > 6) PetscCall(VecGetKokkosView(yp[6], &y6));
    434:     PetscCall(PetscLogGpuTimeBegin());
    435:     switch (rem) {
    436:     case 7:
    437:       Kokkos::parallel_reduce(
    438:         "VecMDot7", policy,
    439:         KOKKOS_LAMBDA(const PetscInt &i, PetscScalar &lsum0, PetscScalar &lsum1, PetscScalar &lsum2, PetscScalar &lsum3, PetscScalar &lsum4, PetscScalar &lsum5, PetscScalar &lsum6) {
    440:         lsum0 += xv(i) * PetscConj(y0(i)); lsum1 += xv(i) * PetscConj(y1(i)); lsum2 += xv(i) * PetscConj(y2(i)); lsum3 += xv(i) * PetscConj(y3(i));
    441:         lsum4 += xv(i) * PetscConj(y4(i)); lsum5 += xv(i) * PetscConj(y5(i)); lsum6 += xv(i) * PetscConj(y6(i));
    442:       }, zp[0], zp[1], zp[2], zp[3], zp[4], zp[5], zp[6]);
    443:       break;
    444:     case 6:
    445:       Kokkos::parallel_reduce(
    446:         "VecMDot6", policy,
    447:         KOKKOS_LAMBDA(const PetscInt &i, PetscScalar &lsum0, PetscScalar &lsum1, PetscScalar &lsum2, PetscScalar &lsum3, PetscScalar &lsum4, PetscScalar &lsum5) {
    448:         lsum0 += xv(i) * PetscConj(y0(i)); lsum1 += xv(i) * PetscConj(y1(i)); lsum2 += xv(i) * PetscConj(y2(i)); lsum3 += xv(i) * PetscConj(y3(i));
    449:         lsum4 += xv(i) * PetscConj(y4(i)); lsum5 += xv(i) * PetscConj(y5(i));
    450:       }, zp[0], zp[1], zp[2], zp[3], zp[4], zp[5]);
    451:       break;
    452:     case 5:
    453:       Kokkos::parallel_reduce(
    454:         "VecMDot5", policy,
    455:         KOKKOS_LAMBDA(const PetscInt &i, PetscScalar &lsum0, PetscScalar &lsum1, PetscScalar &lsum2, PetscScalar &lsum3, PetscScalar &lsum4) {
    456:         lsum0 += xv(i) * PetscConj(y0(i)); lsum1 += xv(i) * PetscConj(y1(i)); lsum2 += xv(i) * PetscConj(y2(i)); lsum3 += xv(i) * PetscConj(y3(i));
    457:         lsum4 += xv(i) * PetscConj(y4(i));
    458:       }, zp[0], zp[1], zp[2], zp[3], zp[4]);
    459:       break;
    460:     case 4:
    461:       Kokkos::parallel_reduce(
    462:         "VecMDot4", policy,
    463:         KOKKOS_LAMBDA(const PetscInt &i, PetscScalar &lsum0, PetscScalar &lsum1, PetscScalar &lsum2, PetscScalar &lsum3) {
    464:         lsum0 += xv(i) * PetscConj(y0(i)); lsum1 += xv(i) * PetscConj(y1(i)); lsum2 += xv(i) * PetscConj(y2(i)); lsum3 += xv(i) * PetscConj(y3(i));
    465:       }, zp[0], zp[1], zp[2], zp[3]);
    466:       break;
    467:     case 3:
    468:       Kokkos::parallel_reduce(
    469:         "VecMDot3", policy,
    470:         KOKKOS_LAMBDA(const PetscInt &i, PetscScalar &lsum0, PetscScalar &lsum1, PetscScalar &lsum2) {
    471:         lsum0 += xv(i) * PetscConj(y0(i)); lsum1 += xv(i) * PetscConj(y1(i)); lsum2 += xv(i) * PetscConj(y2(i));
    472:       }, zp[0], zp[1], zp[2]);
    473:       break;
    474:     case 2:
    475:       Kokkos::parallel_reduce(
    476:         "VecMDot2", policy,
    477:         KOKKOS_LAMBDA(const PetscInt &i, PetscScalar &lsum0, PetscScalar &lsum1) {
    478:         lsum0 += xv(i) * PetscConj(y0(i)); lsum1 += xv(i) * PetscConj(y1(i));
    479:       }, zp[0], zp[1]);
    480:       break;
    481:     case 1:
    482:       Kokkos::parallel_reduce(
    483:         "VecMDot1", policy,
    484:         KOKKOS_LAMBDA(const PetscInt &i, PetscScalar &lsum0) {
    485:         lsum0 += xv(i) * PetscConj(y0(i));
    486:       }, zp[0]);
    487:       break;
    488:     }
    489:     PetscCall(PetscLogGpuTimeEnd());
    490:     PetscCall(PetscLogGpuToCpu(rem * sizeof(PetscScalar))); // for copying to z[] on host
    491:     if (rem > 0) PetscCall(VecRestoreKokkosView(yp[0], &y0));
    492:     if (rem > 1) PetscCall(VecRestoreKokkosView(yp[1], &y1));
    493:     if (rem > 2) PetscCall(VecRestoreKokkosView(yp[2], &y2));
    494:     if (rem > 3) PetscCall(VecRestoreKokkosView(yp[3], &y3));
    495:     if (rem > 4) PetscCall(VecRestoreKokkosView(yp[4], &y4));
    496:     if (rem > 5) PetscCall(VecRestoreKokkosView(yp[5], &y5));
    497:     if (rem > 6) PetscCall(VecRestoreKokkosView(yp[6], &y6));
    498:   }
    499:   PetscCall(VecRestoreKokkosView(xin, &xv));
    500:   PetscFunctionReturn(PETSC_SUCCESS);
    501:   // clang-format on
    502: }

    504: /* z[i] = (x,y_i) = y_i^H x */
    505: PetscErrorCode VecMDot_SeqKokkos(Vec xin, PetscInt nv, const Vec yin[], PetscScalar *z)
    506: {
    507:   PetscFunctionBegin;
    508:   // With no good reason, VecMultiDot_Private() performs much worse than VecMultiDot_Verbose() with HIP,
    509:   // but they are on par with CUDA. Kokkos team is investigating this problem.
    510: #if 0
    511:   PetscCall(VecMultiDot_Private<ConjugateDotTag>(xin, nv, yin, z));
    512: #else
    513:   PetscCall(VecMultiDot_Verbose(xin, nv, yin, z));
    514: #endif
    515:   PetscCall(PetscLogGpuFlops(PetscMax(nv * (2.0 * xin->map->n - 1), 0.0)));
    516:   PetscFunctionReturn(PETSC_SUCCESS);
    517: }

    519: /* z[i] = (x,y_i) = y_i^T x */
    520: PetscErrorCode VecMTDot_SeqKokkos(Vec xin, PetscInt nv, const Vec yin[], PetscScalar *z)
    521: {
    522:   PetscFunctionBegin;
    523:   PetscCall(VecMultiDot_Private<TransposeDotTag>(xin, nv, yin, z));
    524:   PetscCall(PetscLogGpuFlops(PetscMax(nv * (2.0 * xin->map->n - 1), 0.0)));
    525:   PetscFunctionReturn(PETSC_SUCCESS);
    526: }

    528: // z[i] = (x,y_i) = y_i^H x OR y_i^T x
    529: static PetscErrorCode VecMultiDot_SeqKokkos_GEMV(PetscBool conjugate, Vec xin, PetscInt nv, const Vec yin[], PetscScalar *z_h)
    530: {
    531:   PetscInt                   i, j, nfail;
    532:   ConstPetscScalarKokkosView xv, yfirst, ynext;
    533:   const PetscScalar         *yarray;
    534:   PetscBool                  stop  = PETSC_FALSE;
    535:   PetscScalar               *z_d   = nullptr;
    536:   const char                *trans = conjugate ? "C" : "T";
    537:   PetscInt64                 lda   = 0;
    538:   PetscInt                   m, n = xin->map->n;

    540:   PetscFunctionBegin;
    541:   PetscCall(VecGetKokkosView(xin, &xv));
    542: #if defined(KOKKOS_ENABLE_UNIFIED_MEMORY)
    543:   z_d = z_h;
    544: #endif
    545:   i = nfail = 0;
    546:   while (i < nv) {
    547:     // search a sequence of vectors with a fixed stride
    548:     stop = PETSC_FALSE;
    549:     PetscCall(VecGetKokkosView(yin[i], &yfirst));
    550:     yarray = yfirst.data();
    551:     for (j = i + 1; j < nv; j++) {
    552:       PetscCall(VecGetKokkosView(yin[j], &ynext));
    553:       if (j == i + 1) {
    554:         lda = ynext.data() - yarray;                       // arbitrary ptrdiff could be very large
    555:         if (lda < 0 || lda - n > 64) stop = PETSC_TRUE;    // avoid using arbitrary lda; 64 bytes are a big enough alignment in VecDuplicateVecs
    556:       } else if (lda * (j - i) != ynext.data() - yarray) { // not in the same stride? if so, stop searching
    557:         stop = PETSC_TRUE;
    558:       }
    559:       PetscCall(VecRestoreKokkosView(yin[j], &ynext));
    560:       if (stop) break;
    561:     }
    562:     PetscCall(VecRestoreKokkosView(yin[i], &yfirst));

    564:     // found m vectors yin[i..j) with a stride lda at address yarray
    565:     m = j - i;
    566:     if (m > 1) {
    567:       if (!z_d) {
    568:         if (nv > PetscScalarPoolSize) { // rare case
    569:           PetscScalarPoolSize = nv;
    570:           PetscCallCXX(PetscScalarPool = static_cast<PetscScalar *>(Kokkos::kokkos_realloc(PetscScalarPool, PetscScalarPoolSize)));
    571:         }
    572:         z_d = PetscScalarPool;
    573:       }
    574:       const auto &A  = Kokkos::View<const PetscScalar **, Kokkos::LayoutLeft>(yarray, lda, m);
    575:       const auto &Y  = Kokkos::subview(A, std::pair<PetscInt, PetscInt>(0, n), Kokkos::ALL);
    576:       auto        zv = PetscScalarKokkosDualView(PetscScalarKokkosView(z_d + i, m), PetscScalarKokkosViewHost(z_h + i, m));
    577:       PetscCall(PetscLogGpuTimeBegin());
    578:       PetscCallCXX(KokkosBlas::gemv(PetscGetKokkosExecutionSpace(), trans, 1.0, Y, xv, 0.0, zv.view_device()));
    579:       PetscCall(PetscLogGpuTimeEnd());
    580:       PetscCallCXX(zv.modify_device());
    581:       PetscCall(KokkosDualViewSyncHost(zv, PetscGetKokkosExecutionSpace()));
    582:       PetscCall(PetscLogGpuFlops(PetscMax(m * (2.0 * n - 1), 0.0)));
    583:     } else {
    584:       // we only allow falling back on VecDot once, to avoid doing VecMultiDot via individual VecDots
    585:       if (nfail++ == 0) {
    586:         if (conjugate) PetscCall(VecDot_SeqKokkos(xin, yin[i], z_h + i));
    587:         else PetscCall(VecTDot_SeqKokkos(xin, yin[i], z_h + i));
    588:       } else break; // break the while loop
    589:     }
    590:     i = j;
    591:   }
    592:   PetscCall(VecRestoreKokkosView(xin, &xv));
    593:   if (i < nv) { // finish the remaining if any
    594:     if (conjugate) PetscCall(VecMDot_SeqKokkos(xin, nv - i, yin + i, z_h + i));
    595:     else PetscCall(VecMTDot_SeqKokkos(xin, nv - i, yin + i, z_h + i));
    596:   }
    597:   PetscFunctionReturn(PETSC_SUCCESS);
    598: }

    600: PetscErrorCode VecMDot_SeqKokkos_GEMV(Vec xin, PetscInt nv, const Vec yin[], PetscScalar *z)
    601: {
    602:   PetscFunctionBegin;
    603:   PetscCall(VecMultiDot_SeqKokkos_GEMV(PETSC_TRUE, xin, nv, yin, z)); // conjugate
    604:   PetscFunctionReturn(PETSC_SUCCESS);
    605: }

    607: PetscErrorCode VecMTDot_SeqKokkos_GEMV(Vec xin, PetscInt nv, const Vec yin[], PetscScalar *z)
    608: {
    609:   PetscFunctionBegin;
    610:   PetscCall(VecMultiDot_SeqKokkos_GEMV(PETSC_FALSE, xin, nv, yin, z)); // transpose
    611:   PetscFunctionReturn(PETSC_SUCCESS);
    612: }

    614: /* x[:] = alpha */
    615: PetscErrorCode VecSet_SeqKokkos(Vec xin, PetscScalar alpha)
    616: {
    617:   PetscScalarKokkosView xv;
    618:   auto                  exec = PetscGetKokkosExecutionSpace();

    620:   PetscFunctionBegin;
    621:   PetscCall(PetscLogGpuTimeBegin());
    622:   PetscCall(VecGetKokkosViewWrite(xin, &xv));
    623:   PetscCallCXX(KokkosBlas::fill(exec, xv, alpha));
    624:   PetscCall(VecRestoreKokkosViewWrite(xin, &xv));
    625:   PetscCall(PetscLogGpuTimeEnd());
    626:   PetscFunctionReturn(PETSC_SUCCESS);
    627: }

    629: /* x = alpha x */
    630: PetscErrorCode VecScale_SeqKokkos(Vec xin, PetscScalar alpha)
    631: {
    632:   auto exec = PetscGetKokkosExecutionSpace();

    634:   PetscFunctionBegin;
    635:   if (alpha == (PetscScalar)0.0) {
    636:     PetscCall(VecSet_SeqKokkos(xin, alpha));
    637:   } else if (alpha != (PetscScalar)1.0) {
    638:     PetscScalarKokkosView xv;

    640:     PetscCall(PetscLogGpuTimeBegin());
    641:     PetscCall(VecGetKokkosView(xin, &xv));
    642:     PetscCallCXX(KokkosBlas::scal(exec, xv, alpha, xv));
    643:     PetscCall(VecRestoreKokkosView(xin, &xv));
    644:     PetscCall(PetscLogGpuTimeEnd());
    645:     PetscCall(PetscLogGpuFlops(xin->map->n));
    646:   }
    647:   PetscFunctionReturn(PETSC_SUCCESS);
    648: }

    650: /* z = y^H x */
    651: PetscErrorCode VecDot_SeqKokkos(Vec xin, Vec yin, PetscScalar *z)
    652: {
    653:   ConstPetscScalarKokkosView xv, yv;
    654:   auto                       exec = PetscGetKokkosExecutionSpace();

    656:   PetscFunctionBegin;
    657:   PetscCall(PetscLogGpuTimeBegin());
    658:   PetscCall(VecGetKokkosView(xin, &xv));
    659:   PetscCall(VecGetKokkosView(yin, &yv));
    660:   PetscCallCXX(*z = KokkosBlas::dot(exec, yv, xv)); /* KokkosBlas::dot(a,b) takes conjugate of a */
    661:   PetscCall(VecRestoreKokkosView(xin, &xv));
    662:   PetscCall(VecRestoreKokkosView(yin, &yv));
    663:   PetscCall(PetscLogGpuTimeEnd());
    664:   PetscCall(PetscLogGpuFlops(PetscMax(2.0 * xin->map->n - 1, 0.0)));
    665:   PetscFunctionReturn(PETSC_SUCCESS);
    666: }

    668: /* y = x, where x is VECKOKKOS, but y may be not */
    669: PetscErrorCode VecCopy_SeqKokkos(Vec xin, Vec yin)
    670: {
    671:   auto exec = PetscGetKokkosExecutionSpace();

    673:   PetscFunctionBegin;
    674:   PetscCall(PetscLogGpuTimeBegin());
    675:   if (xin != yin) {
    676:     Vec_Kokkos *xkok = static_cast<Vec_Kokkos *>(xin->spptr);
    677:     if (yin->offloadmask == PETSC_OFFLOAD_KOKKOS) {
    678:       /* y is also a VecKokkos */
    679:       Vec_Kokkos *ykok = static_cast<Vec_Kokkos *>(yin->spptr);
    680:       /* Kokkos rule: if x's host has newer data, it will copy to y's host view; otherwise to y's device view
    681:         In case x's host is newer, y's device is newer, it will error (though should not, I think). So we just
    682:         clear y's sync state.
    683:        */
    684:       ykok->v_dual.clear_sync_state();
    685:       PetscCallCXX(Kokkos::deep_copy(exec, ykok->v_dual, xkok->v_dual)); // either cpu2cpu or gpu2cpu, so don't log it
    686:     } else {
    687:       PetscScalar *yarray;
    688:       PetscCall(VecGetArrayWrite(yin, &yarray));
    689:       PetscScalarKokkosViewHost yv(yarray, yin->map->n);
    690:       if (xkok->v_dual.need_sync_host()) {                                     // x's device has newer data
    691:         PetscCallCXX(Kokkos::deep_copy(exec, yv, xkok->v_dual.view_device())); // gpu2cpu
    692:         PetscCallCXX(exec.fence());                                            // finish the deep copy
    693:         PetscCall(PetscLogGpuToCpu(xkok->v_dual.extent(0) * sizeof(PetscScalar)));
    694:       } else {
    695:         PetscCallCXX(exec.fence());                                          // make sure xkok->v_dual.view_host() in ready for use on host;  Kokkos might also call it inside deep_copy(). We do it here for safety.
    696:         PetscCallCXX(Kokkos::deep_copy(exec, yv, xkok->v_dual.view_host())); // Host view to host view deep copy, done on host
    697:       }
    698:       PetscCall(VecRestoreArrayWrite(yin, &yarray));
    699:     }
    700:   }
    701:   PetscCall(PetscLogGpuTimeEnd());
    702:   PetscFunctionReturn(PETSC_SUCCESS);
    703: }

    705: /* y[i] <--> x[i] */
    706: PetscErrorCode VecSwap_SeqKokkos(Vec xin, Vec yin)
    707: {
    708:   PetscFunctionBegin;
    709:   if (xin != yin) {
    710:     PetscScalarKokkosView xv, yv;

    712:     PetscCall(PetscLogGpuTimeBegin());
    713:     PetscCall(VecGetKokkosView(xin, &xv));
    714:     PetscCall(VecGetKokkosView(yin, &yv));
    715:     PetscCallCXX(Kokkos::parallel_for(
    716:       Kokkos::RangePolicy<>(PetscGetKokkosExecutionSpace(), 0, xin->map->n), KOKKOS_LAMBDA(const PetscInt &i) {
    717:         PetscScalar tmp = xv(i);
    718:         xv(i)           = yv(i);
    719:         yv(i)           = tmp;
    720:       }));
    721:     PetscCall(VecRestoreKokkosView(xin, &xv));
    722:     PetscCall(VecRestoreKokkosView(yin, &yv));
    723:     PetscCall(PetscLogGpuTimeEnd());
    724:   }
    725:   PetscFunctionReturn(PETSC_SUCCESS);
    726: }

    728: /*  w = alpha x + y */
    729: PetscErrorCode VecWAXPY_SeqKokkos(Vec win, PetscScalar alpha, Vec xin, Vec yin)
    730: {
    731:   PetscFunctionBegin;
    732:   if (alpha == (PetscScalar)0.0) {
    733:     PetscCall(VecCopy_SeqKokkos(yin, win));
    734:   } else {
    735:     ConstPetscScalarKokkosView xv, yv;
    736:     PetscScalarKokkosView      wv;

    738:     PetscCall(PetscLogGpuTimeBegin());
    739:     PetscCall(VecGetKokkosViewWrite(win, &wv));
    740:     PetscCall(VecGetKokkosView(xin, &xv));
    741:     PetscCall(VecGetKokkosView(yin, &yv));
    742:     PetscCallCXX(Kokkos::parallel_for(Kokkos::RangePolicy<>(PetscGetKokkosExecutionSpace(), 0, win->map->n), KOKKOS_LAMBDA(const PetscInt &i) { wv(i) = alpha * xv(i) + yv(i); }));
    743:     PetscCall(VecRestoreKokkosView(xin, &xv));
    744:     PetscCall(VecRestoreKokkosView(yin, &yv));
    745:     PetscCall(VecRestoreKokkosViewWrite(win, &wv));
    746:     PetscCall(PetscLogGpuTimeEnd());
    747:     PetscCall(PetscLogGpuFlops(2.0 * win->map->n));
    748:   }
    749:   PetscFunctionReturn(PETSC_SUCCESS);
    750: }

    752: template <PetscInt ValueCount>
    753: struct MAXPYFunctor {
    754:   static_assert(ValueCount >= 1 && ValueCount <= 8, "ValueCount must be in [1, 8]");
    755:   typedef ConstPetscScalarKokkosView::size_type size_type;

    757:   PetscScalarKokkosView      yv;
    758:   PetscScalar                a[8];
    759:   ConstPetscScalarKokkosView xv[8];

    761:   MAXPYFunctor(PetscScalarKokkosView yv, PetscScalar a0, PetscScalar a1, PetscScalar a2, PetscScalar a3, PetscScalar a4, PetscScalar a5, PetscScalar a6, PetscScalar a7, ConstPetscScalarKokkosView xv0, ConstPetscScalarKokkosView xv1, ConstPetscScalarKokkosView xv2, ConstPetscScalarKokkosView xv3, ConstPetscScalarKokkosView xv4, ConstPetscScalarKokkosView xv5, ConstPetscScalarKokkosView xv6, ConstPetscScalarKokkosView xv7) :
    762:     yv(yv)
    763:   {
    764:     a[0]  = a0;
    765:     a[1]  = a1;
    766:     a[2]  = a2;
    767:     a[3]  = a3;
    768:     a[4]  = a4;
    769:     a[5]  = a5;
    770:     a[6]  = a6;
    771:     a[7]  = a7;
    772:     xv[0] = xv0;
    773:     xv[1] = xv1;
    774:     xv[2] = xv2;
    775:     xv[3] = xv3;
    776:     xv[4] = xv4;
    777:     xv[5] = xv5;
    778:     xv[6] = xv6;
    779:     xv[7] = xv7;
    780:   }

    782:   KOKKOS_INLINE_FUNCTION void operator()(const size_type i) const
    783:   {
    784:     for (PetscInt j = 0; j < ValueCount; ++j) yv(i) += a[j] * xv[j](i);
    785:   }
    786: };

    788: /*  y = y + sum alpha[i] x[i] */
    789: PetscErrorCode VecMAXPY_SeqKokkos(Vec yin, PetscInt nv, const PetscScalar *alpha, Vec *xin)
    790: {
    791:   PetscInt                   i, j, cur = 0, ngroup = nv / 8, rem = nv % 8, N = yin->map->n;
    792:   PetscScalarKokkosView      yv;
    793:   PetscScalar                a[8];
    794:   ConstPetscScalarKokkosView xv[8];
    795:   Kokkos::RangePolicy<>      policy(PetscGetKokkosExecutionSpace(), 0, N);

    797:   PetscFunctionBegin;
    798:   PetscCall(PetscLogGpuTimeBegin());
    799:   PetscCall(VecGetKokkosView(yin, &yv));
    800:   for (i = 0; i < ngroup; i++) { /* 8 x's per group */
    801:     for (j = 0; j < 8; j++) {    /* Fill the parameters */
    802:       a[j] = alpha[cur + j];
    803:       PetscCall(VecGetKokkosView(xin[cur + j], &xv[j]));
    804:     }
    805:     MAXPYFunctor<8> maxpy(yv, a[0], a[1], a[2], a[3], a[4], a[5], a[6], a[7], xv[0], xv[1], xv[2], xv[3], xv[4], xv[5], xv[6], xv[7]);
    806:     PetscCallCXX(Kokkos::parallel_for(policy, maxpy));
    807:     for (j = 0; j < 8; j++) PetscCall(VecRestoreKokkosView(xin[cur + j], &xv[j]));
    808:     cur += 8;
    809:   }

    811:   if (rem) { /* The remaining */
    812:     for (j = 0; j < rem; j++) {
    813:       a[j] = alpha[cur + j];
    814:       PetscCall(VecGetKokkosView(xin[cur + j], &xv[j]));
    815:     }
    816:     // clang-format off
    817:     switch (rem) {
    818:     case 1: PetscCallCXX(Kokkos::parallel_for(policy, MAXPYFunctor<1>(yv, a[0], a[1], a[2], a[3], a[4], a[5], a[6], a[7], xv[0], xv[1], xv[2], xv[3], xv[4], xv[5], xv[6], xv[7]))); break;
    819:     case 2: PetscCallCXX(Kokkos::parallel_for(policy, MAXPYFunctor<2>(yv, a[0], a[1], a[2], a[3], a[4], a[5], a[6], a[7], xv[0], xv[1], xv[2], xv[3], xv[4], xv[5], xv[6], xv[7]))); break;
    820:     case 3: PetscCallCXX(Kokkos::parallel_for(policy, MAXPYFunctor<3>(yv, a[0], a[1], a[2], a[3], a[4], a[5], a[6], a[7], xv[0], xv[1], xv[2], xv[3], xv[4], xv[5], xv[6], xv[7]))); break;
    821:     case 4: PetscCallCXX(Kokkos::parallel_for(policy, MAXPYFunctor<4>(yv, a[0], a[1], a[2], a[3], a[4], a[5], a[6], a[7], xv[0], xv[1], xv[2], xv[3], xv[4], xv[5], xv[6], xv[7]))); break;
    822:     case 5: PetscCallCXX(Kokkos::parallel_for(policy, MAXPYFunctor<5>(yv, a[0], a[1], a[2], a[3], a[4], a[5], a[6], a[7], xv[0], xv[1], xv[2], xv[3], xv[4], xv[5], xv[6], xv[7]))); break;
    823:     case 6: PetscCallCXX(Kokkos::parallel_for(policy, MAXPYFunctor<6>(yv, a[0], a[1], a[2], a[3], a[4], a[5], a[6], a[7], xv[0], xv[1], xv[2], xv[3], xv[4], xv[5], xv[6], xv[7]))); break;
    824:     case 7: PetscCallCXX(Kokkos::parallel_for(policy, MAXPYFunctor<7>(yv, a[0], a[1], a[2], a[3], a[4], a[5], a[6], a[7], xv[0], xv[1], xv[2], xv[3], xv[4], xv[5], xv[6], xv[7]))); break;
    825:     }
    826:     // clang-format on
    827:     for (j = 0; j < rem; j++) PetscCall(VecRestoreKokkosView(xin[cur + j], &xv[j]));
    828:   }
    829:   PetscCall(VecRestoreKokkosView(yin, &yv));
    830:   PetscCall(PetscLogGpuTimeEnd());
    831:   PetscCall(PetscLogGpuFlops(nv * 2.0 * yin->map->n));
    832:   PetscFunctionReturn(PETSC_SUCCESS);
    833: }

    835: /*  y = y + sum alpha[i] x[i] */
    836: PetscErrorCode VecMAXPY_SeqKokkos_GEMV(Vec yin, PetscInt nv, const PetscScalar *a_h, Vec *xin)
    837: {
    838:   const PetscInt             n = yin->map->n;
    839:   PetscInt                   i, j, nfail;
    840:   PetscScalarKokkosView      yv;
    841:   ConstPetscScalarKokkosView xfirst, xnext;
    842:   PetscBool                  stop = PETSC_FALSE;
    843:   PetscInt                   lda  = 0, m;
    844:   const PetscScalar         *xarray;
    845:   PetscScalar               *a_d = nullptr;

    847:   PetscFunctionBegin;
    848:   PetscCall(PetscLogGpuTimeBegin());
    849:   PetscCall(VecGetKokkosView(yin, &yv));
    850: #if defined(KOKKOS_ENABLE_UNIFIED_MEMORY)
    851:   a_d = const_cast<PetscScalar *>(a_h);
    852: #endif
    853:   i = nfail = 0;
    854:   while (i < nv) {
    855:     stop = PETSC_FALSE;
    856:     PetscCall(VecGetKokkosView(xin[i], &xfirst));
    857:     xarray = xfirst.data();
    858:     for (j = i + 1; j < nv; j++) {
    859:       PetscCall(VecGetKokkosView(xin[j], &xnext));
    860:       if (j == i + 1) {
    861:         lda = xnext.data() - xfirst.data();
    862:         if (lda < 0 || lda - n > 64) stop = PETSC_TRUE;    // avoid using arbitrary lda; 64 bytes are a big enough alignment in VecDuplicateVecs
    863:       } else if (lda * (j - i) != xnext.data() - xarray) { // not in the same stride? if so, stop here
    864:         stop = PETSC_TRUE;
    865:       }
    866:       PetscCall(VecRestoreKokkosView(xin[j], &xnext));
    867:       if (stop) break;
    868:     }
    869:     PetscCall(VecRestoreKokkosView(xin[i], &xfirst));

    871:     m = j - i;
    872:     if (m > 1) {
    873:       if (!a_d) {
    874:         if (nv > PetscScalarPoolSize) { // rare case
    875:           PetscScalarPoolSize = nv;
    876:           PetscCallCXX(PetscScalarPool = static_cast<PetscScalar *>(Kokkos::kokkos_realloc(PetscScalarPool, PetscScalarPoolSize)));
    877:         }
    878:         a_d = PetscScalarPool;
    879:       }
    880:       const auto &B  = Kokkos::View<const PetscScalar **, Kokkos::LayoutLeft>(xarray, lda, m);
    881:       const auto &A  = Kokkos::subview(B, std::pair<PetscInt, PetscInt>(0, n), Kokkos::ALL);
    882:       auto        av = PetscScalarKokkosDualView(PetscScalarKokkosView(a_d + i, m), PetscScalarKokkosViewHost(const_cast<PetscScalar *>(a_h) + i, m));
    883:       av.modify_host();
    884:       PetscCall(KokkosDualViewSyncDevice(av, PetscGetKokkosExecutionSpace()));
    885:       PetscCallCXX(KokkosBlas::gemv(PetscGetKokkosExecutionSpace(), "N", 1.0, A, av.view_device(), 1.0, yv));
    886:       PetscCall(PetscLogGpuFlops(m * 2.0 * n));
    887:     } else {
    888:       // we only allow falling back on VecAXPY once
    889:       if (nfail++ == 0) PetscCall(VecAXPY_SeqKokkos(yin, a_h[i], xin[i]));
    890:       else break; // break the while loop
    891:     }
    892:     i = j;
    893:   }
    894:   // finish the remaining if any
    895:   PetscCall(VecRestoreKokkosView(yin, &yv));
    896:   if (i < nv) PetscCall(VecMAXPY_SeqKokkos(yin, nv - i, a_h + i, xin + i));
    897:   PetscCall(PetscLogGpuTimeEnd());
    898:   PetscFunctionReturn(PETSC_SUCCESS);
    899: }

    901: /* y = alpha x + beta y */
    902: PetscErrorCode VecAXPBY_SeqKokkos(Vec yin, PetscScalar alpha, PetscScalar beta, Vec xin)
    903: {
    904:   PetscBool xiskok, yiskok;

    906:   PetscFunctionBegin;
    907:   PetscCall(PetscObjectTypeCompareAny((PetscObject)xin, &xiskok, VECSEQKOKKOS, VECMPIKOKKOS, ""));
    908:   PetscCall(PetscObjectTypeCompareAny((PetscObject)yin, &yiskok, VECSEQKOKKOS, VECMPIKOKKOS, ""));
    909:   if (xiskok && yiskok) {
    910:     ConstPetscScalarKokkosView xv;
    911:     PetscScalarKokkosView      yv;

    913:     PetscCall(PetscLogGpuTimeBegin());
    914:     PetscCall(VecGetKokkosView(xin, &xv));
    915:     PetscCall(VecGetKokkosView(yin, &yv));
    916:     PetscCallCXX(KokkosBlas::axpby(PetscGetKokkosExecutionSpace(), alpha, xv, beta, yv));
    917:     PetscCall(VecRestoreKokkosView(xin, &xv));
    918:     PetscCall(VecRestoreKokkosView(yin, &yv));
    919:     PetscCall(PetscLogGpuTimeEnd());
    920:     if (alpha == (PetscScalar)0.0 || beta == (PetscScalar)0.0) {
    921:       PetscCall(PetscLogGpuFlops(xin->map->n));
    922:     } else if (beta == (PetscScalar)1.0 || alpha == (PetscScalar)1.0) {
    923:       PetscCall(PetscLogGpuFlops(2.0 * xin->map->n));
    924:     } else {
    925:       PetscCall(PetscLogGpuFlops(3.0 * xin->map->n));
    926:     }
    927:   } else {
    928:     PetscCall(VecAXPBY_Seq(yin, alpha, beta, xin));
    929:   }
    930:   PetscFunctionReturn(PETSC_SUCCESS);
    931: }

    933: /* z = alpha x + beta y + gamma z */
    934: PetscErrorCode VecAXPBYPCZ_SeqKokkos(Vec zin, PetscScalar alpha, PetscScalar beta, PetscScalar gamma, Vec xin, Vec yin)
    935: {
    936:   ConstPetscScalarKokkosView xv, yv;
    937:   PetscScalarKokkosView      zv;
    938:   Kokkos::RangePolicy<>      policy(PetscGetKokkosExecutionSpace(), 0, zin->map->n);

    940:   PetscFunctionBegin;
    941:   PetscCall(PetscLogGpuTimeBegin());
    942:   PetscCall(VecGetKokkosView(zin, &zv));
    943:   PetscCall(VecGetKokkosView(xin, &xv));
    944:   PetscCall(VecGetKokkosView(yin, &yv));
    945:   if (gamma == (PetscScalar)0.0) { // a common case
    946:     if (alpha == -beta) {
    947:       PetscCallCXX(Kokkos::parallel_for( // a common case
    948:         policy, KOKKOS_LAMBDA(const PetscInt &i) { zv(i) = alpha * (xv(i) - yv(i)); }));
    949:     } else {
    950:       PetscCallCXX(Kokkos::parallel_for(policy, KOKKOS_LAMBDA(const PetscInt &i) { zv(i) = alpha * xv(i) + beta * yv(i); }));
    951:     }
    952:   } else {
    953:     PetscCallCXX(KokkosBlas::update(PetscGetKokkosExecutionSpace(), alpha, xv, beta, yv, gamma, zv));
    954:   }
    955:   PetscCall(VecRestoreKokkosView(xin, &xv));
    956:   PetscCall(VecRestoreKokkosView(yin, &yv));
    957:   PetscCall(VecRestoreKokkosView(zin, &zv));
    958:   PetscCall(PetscLogGpuTimeEnd());
    959:   PetscCall(PetscLogGpuFlops(zin->map->n * 5.0));
    960:   PetscFunctionReturn(PETSC_SUCCESS);
    961: }

    963: /* w = x*y. Any subset of the x, y, and w may be the same vector.

    965:   w is of type VecKokkos, but x, y may be not.
    966: */
    967: PetscErrorCode VecPointwiseMult_SeqKokkos(Vec win, Vec xin, Vec yin)
    968: {
    969:   PetscInt n;

    971:   PetscFunctionBegin;
    972:   PetscCall(PetscLogGpuTimeBegin());
    973:   PetscCall(VecGetLocalSize(win, &n));
    974:   if (xin->offloadmask != PETSC_OFFLOAD_KOKKOS || yin->offloadmask != PETSC_OFFLOAD_KOKKOS) {
    975:     PetscScalarKokkosViewHost wv;
    976:     const PetscScalar        *xp, *yp;
    977:     PetscCall(VecGetArrayRead(xin, &xp));
    978:     PetscCall(VecGetArrayRead(yin, &yp));
    979:     PetscCall(VecGetKokkosViewWrite(win, &wv));

    981:     ConstPetscScalarKokkosViewHost xv(xp, n), yv(yp, n);
    982:     PetscCallCXX(Kokkos::parallel_for(Kokkos::RangePolicy<Kokkos::DefaultHostExecutionSpace>(0, n), KOKKOS_LAMBDA(const PetscInt &i) { wv(i) = xv(i) * yv(i); }));

    984:     PetscCall(VecRestoreArrayRead(xin, &xp));
    985:     PetscCall(VecRestoreArrayRead(yin, &yp));
    986:     PetscCall(VecRestoreKokkosViewWrite(win, &wv));
    987:   } else {
    988:     ConstPetscScalarKokkosView xv, yv;
    989:     PetscScalarKokkosView      wv;

    991:     PetscCall(VecGetKokkosViewWrite(win, &wv));
    992:     PetscCall(VecGetKokkosView(xin, &xv));
    993:     PetscCall(VecGetKokkosView(yin, &yv));
    994:     PetscCallCXX(Kokkos::parallel_for(Kokkos::RangePolicy<>(PetscGetKokkosExecutionSpace(), 0, n), KOKKOS_LAMBDA(const PetscInt &i) { wv(i) = xv(i) * yv(i); }));
    995:     PetscCall(VecRestoreKokkosView(yin, &yv));
    996:     PetscCall(VecRestoreKokkosView(xin, &xv));
    997:     PetscCall(VecRestoreKokkosViewWrite(win, &wv));
    998:   }
    999:   PetscCall(PetscLogGpuTimeEnd());
    1000:   PetscCall(PetscLogGpuFlops(n));
    1001:   PetscFunctionReturn(PETSC_SUCCESS);
    1002: }

    1004: /* w = x/y */
    1005: PetscErrorCode VecPointwiseDivide_SeqKokkos(Vec win, Vec xin, Vec yin)
    1006: {
    1007:   PetscInt n;

    1009:   PetscFunctionBegin;
    1010:   PetscCall(PetscLogGpuTimeBegin());
    1011:   PetscCall(VecGetLocalSize(win, &n));
    1012:   if (xin->offloadmask != PETSC_OFFLOAD_KOKKOS || yin->offloadmask != PETSC_OFFLOAD_KOKKOS) {
    1013:     PetscScalarKokkosViewHost wv;
    1014:     const PetscScalar        *xp, *yp;
    1015:     PetscCall(VecGetArrayRead(xin, &xp));
    1016:     PetscCall(VecGetArrayRead(yin, &yp));
    1017:     PetscCall(VecGetKokkosViewWrite(win, &wv));

    1019:     ConstPetscScalarKokkosViewHost xv(xp, n), yv(yp, n);
    1020:     PetscCallCXX(Kokkos::parallel_for(
    1021:       Kokkos::RangePolicy<Kokkos::DefaultHostExecutionSpace>(0, n), KOKKOS_LAMBDA(const PetscInt &i) {
    1022:         if (yv(i) != 0.0) wv(i) = xv(i) / yv(i);
    1023:         else wv(i) = 0.0;
    1024:       }));

    1026:     PetscCall(VecRestoreArrayRead(xin, &xp));
    1027:     PetscCall(VecRestoreArrayRead(yin, &yp));
    1028:     PetscCall(VecRestoreKokkosViewWrite(win, &wv));
    1029:   } else {
    1030:     ConstPetscScalarKokkosView xv, yv;
    1031:     PetscScalarKokkosView      wv;

    1033:     PetscCall(VecGetKokkosViewWrite(win, &wv));
    1034:     PetscCall(VecGetKokkosView(xin, &xv));
    1035:     PetscCall(VecGetKokkosView(yin, &yv));
    1036:     PetscCallCXX(Kokkos::parallel_for(
    1037:       Kokkos::RangePolicy<>(PetscGetKokkosExecutionSpace(), 0, n), KOKKOS_LAMBDA(const PetscInt &i) {
    1038:         if (yv(i) != 0.0) wv(i) = xv(i) / yv(i);
    1039:         else wv(i) = 0.0;
    1040:       }));
    1041:     PetscCall(VecRestoreKokkosView(yin, &yv));
    1042:     PetscCall(VecRestoreKokkosView(xin, &xv));
    1043:     PetscCall(VecRestoreKokkosViewWrite(win, &wv));
    1044:   }
    1045:   PetscCall(PetscLogGpuTimeEnd());
    1046:   PetscCall(PetscLogGpuFlops(win->map->n));
    1047:   PetscFunctionReturn(PETSC_SUCCESS);
    1048: }

    1050: PetscErrorCode VecNorm_SeqKokkos(Vec xin, NormType type, PetscReal *z)
    1051: {
    1052:   const PetscInt             n = xin->map->n;
    1053:   ConstPetscScalarKokkosView xv;
    1054:   auto                       exec = PetscGetKokkosExecutionSpace();

    1056:   PetscFunctionBegin;
    1057:   if (type == NORM_1_AND_2) {
    1058:     PetscCall(VecNorm_SeqKokkos(xin, NORM_1, z));
    1059:     PetscCall(VecNorm_SeqKokkos(xin, NORM_2, z + 1));
    1060:   } else {
    1061:     PetscCall(PetscLogGpuTimeBegin());
    1062:     PetscCall(VecGetKokkosView(xin, &xv));
    1063:     if (type == NORM_2 || type == NORM_FROBENIUS) {
    1064:       PetscCallCXX(*z = KokkosBlas::nrm2(exec, xv));
    1065:       PetscCall(PetscLogGpuFlops(PetscMax(2.0 * n - 1, 0.0)));
    1066:     } else if (type == NORM_1) {
    1067:       PetscCallCXX(*z = KokkosBlas::nrm1(exec, xv));
    1068:       PetscCall(PetscLogGpuFlops(PetscMax(n - 1.0, 0.0)));
    1069:     } else if (type == NORM_INFINITY) {
    1070:       PetscCallCXX(*z = KokkosBlas::nrminf(exec, xv));
    1071:     }
    1072:     PetscCall(VecRestoreKokkosView(xin, &xv));
    1073:     PetscCall(PetscLogGpuTimeEnd());
    1074:   }
    1075:   PetscFunctionReturn(PETSC_SUCCESS);
    1076: }

    1078: PetscErrorCode VecErrorWeightedNorms_SeqKokkos(Vec U, Vec Y, Vec E, NormType wnormtype, PetscReal atol, Vec vatol, PetscReal rtol, Vec vrtol, PetscReal ignore_max, PetscReal *norm, PetscInt *norm_loc, PetscReal *norma, PetscInt *norma_loc, PetscReal *normr, PetscInt *normr_loc)
    1079: {
    1080:   ConstPetscScalarKokkosView u, y, erra, atola, rtola;
    1081:   PetscBool                  has_E = PETSC_FALSE, has_atol = PETSC_FALSE, has_rtol = PETSC_FALSE;
    1082:   PetscInt                   n, n_loc = 0, na_loc = 0, nr_loc = 0;
    1083:   PetscReal                  nrm = 0, nrma = 0, nrmr = 0;

    1085:   PetscFunctionBegin;
    1086:   PetscCall(VecGetLocalSize(U, &n));
    1087:   PetscCall(VecGetKokkosView(U, &u));
    1088:   PetscCall(VecGetKokkosView(Y, &y));
    1089:   if (E) {
    1090:     PetscCall(VecGetKokkosView(E, &erra));
    1091:     has_E = PETSC_TRUE;
    1092:   }
    1093:   if (vatol) {
    1094:     PetscCall(VecGetKokkosView(vatol, &atola));
    1095:     has_atol = PETSC_TRUE;
    1096:   }
    1097:   if (vrtol) {
    1098:     PetscCall(VecGetKokkosView(vrtol, &rtola));
    1099:     has_rtol = PETSC_TRUE;
    1100:   }

    1102:   if (wnormtype == NORM_INFINITY) {
    1103:     PetscCallCXX(Kokkos::parallel_reduce(
    1104:       "VecErrorWeightedNorms_INFINITY", Kokkos::RangePolicy<>(PetscGetKokkosExecutionSpace(), 0, n),
    1105:       KOKKOS_LAMBDA(const PetscInt &i, PetscReal &l_nrm, PetscReal &l_nrma, PetscReal &l_nrmr, PetscInt &l_n_loc, PetscInt &l_na_loc, PetscInt &l_nr_loc) {
    1106:         PetscReal err, tol, tola, tolr, l_atol, l_rtol;
    1107:         if (PetscAbsScalar(y(i)) >= ignore_max && PetscAbsScalar(u(i)) >= ignore_max) {
    1108:           l_atol = has_atol ? PetscRealPart(atola(i)) : atol;
    1109:           l_rtol = has_rtol ? PetscRealPart(rtola(i)) : rtol;
    1110:           err    = has_E ? PetscAbsScalar(erra(i)) : PetscAbsScalar(y(i) - u(i));
    1111:           tola   = l_atol;
    1112:           tolr   = l_rtol * PetscMax(PetscAbsScalar(u(i)), PetscAbsScalar(y(i)));
    1113:           tol    = tola + tolr;
    1114:           if (tola > 0.) {
    1115:             l_nrma = PetscMax(l_nrma, err / tola);
    1116:             l_na_loc++;
    1117:           }
    1118:           if (tolr > 0.) {
    1119:             l_nrmr = PetscMax(l_nrmr, err / tolr);
    1120:             l_nr_loc++;
    1121:           }
    1122:           if (tol > 0.) {
    1123:             l_nrm = PetscMax(l_nrm, err / tol);
    1124:             l_n_loc++;
    1125:           }
    1126:         }
    1127:       },
    1128:       Kokkos::Max<PetscReal>(nrm), Kokkos::Max<PetscReal>(nrma), Kokkos::Max<PetscReal>(nrmr), n_loc, na_loc, nr_loc));
    1129:   } else {
    1130:     PetscCallCXX(Kokkos::parallel_reduce(
    1131:       "VecErrorWeightedNorms_NORM_2", Kokkos::RangePolicy<>(PetscGetKokkosExecutionSpace(), 0, n),
    1132:       KOKKOS_LAMBDA(const PetscInt &i, PetscReal &l_nrm, PetscReal &l_nrma, PetscReal &l_nrmr, PetscInt &l_n_loc, PetscInt &l_na_loc, PetscInt &l_nr_loc) {
    1133:         PetscReal err, tol, tola, tolr, l_atol, l_rtol;
    1134:         if (PetscAbsScalar(y(i)) >= ignore_max && PetscAbsScalar(u(i)) >= ignore_max) {
    1135:           l_atol = has_atol ? PetscRealPart(atola(i)) : atol;
    1136:           l_rtol = has_rtol ? PetscRealPart(rtola(i)) : rtol;
    1137:           err    = has_E ? PetscAbsScalar(erra(i)) : PetscAbsScalar(y(i) - u(i));
    1138:           tola   = l_atol;
    1139:           tolr   = l_rtol * PetscMax(PetscAbsScalar(u(i)), PetscAbsScalar(y(i)));
    1140:           tol    = tola + tolr;
    1141:           if (tola > 0.) {
    1142:             l_nrma += PetscSqr(err / tola);
    1143:             l_na_loc++;
    1144:           }
    1145:           if (tolr > 0.) {
    1146:             l_nrmr += PetscSqr(err / tolr);
    1147:             l_nr_loc++;
    1148:           }
    1149:           if (tol > 0.) {
    1150:             l_nrm += PetscSqr(err / tol);
    1151:             l_n_loc++;
    1152:           }
    1153:         }
    1154:       },
    1155:       nrm, nrma, nrmr, n_loc, na_loc, nr_loc));
    1156:   }

    1158:   if (wnormtype == NORM_2) {
    1159:     *norm  = PetscSqrtReal(nrm);
    1160:     *norma = PetscSqrtReal(nrma);
    1161:     *normr = PetscSqrtReal(nrmr);
    1162:   } else {
    1163:     *norm  = nrm;
    1164:     *norma = nrma;
    1165:     *normr = nrmr;
    1166:   }
    1167:   *norm_loc  = n_loc;
    1168:   *norma_loc = na_loc;
    1169:   *normr_loc = nr_loc;

    1171:   if (E) PetscCall(VecRestoreKokkosView(E, &erra));
    1172:   if (vatol) PetscCall(VecRestoreKokkosView(vatol, &atola));
    1173:   if (vrtol) PetscCall(VecRestoreKokkosView(vrtol, &rtola));
    1174:   PetscCall(VecRestoreKokkosView(U, &u));
    1175:   PetscCall(VecRestoreKokkosView(Y, &y));
    1176:   PetscFunctionReturn(PETSC_SUCCESS);
    1177: }

    1179: /* A functor for DotNorm2 so that we can compute dp and nm in one kernel */
    1180: struct DotNorm2 {
    1181:   typedef PetscScalar                           value_type[];
    1182:   typedef ConstPetscScalarKokkosView::size_type size_type;

    1184:   size_type                  value_count;
    1185:   ConstPetscScalarKokkosView xv_, yv_; /* first and second vectors in VecDotNorm2. The order matters. */

    1187:   DotNorm2(ConstPetscScalarKokkosView &xv, ConstPetscScalarKokkosView &yv) : value_count(2), xv_(xv), yv_(yv) { }

    1189:   KOKKOS_INLINE_FUNCTION void operator()(const size_type i, value_type result) const
    1190:   {
    1191:     result[0] += PetscConj(yv_(i)) * xv_(i);
    1192:     result[1] += PetscConj(yv_(i)) * yv_(i);
    1193:   }

    1195:   KOKKOS_INLINE_FUNCTION void join(value_type dst, const value_type src) const
    1196:   {
    1197:     dst[0] += src[0];
    1198:     dst[1] += src[1];
    1199:   }

    1201:   KOKKOS_INLINE_FUNCTION void init(value_type result) const
    1202:   {
    1203:     result[0] = 0.0;
    1204:     result[1] = 0.0;
    1205:   }
    1206: };

    1208: /* dp = y^H x, nm = y^H y */
    1209: PetscErrorCode VecDotNorm2_SeqKokkos(Vec xin, Vec yin, PetscScalar *dp, PetscScalar *nm)
    1210: {
    1211:   ConstPetscScalarKokkosView xv, yv;
    1212:   PetscScalar                result[2];

    1214:   PetscFunctionBegin;
    1215:   PetscCall(PetscLogGpuTimeBegin());
    1216:   PetscCall(VecGetKokkosView(xin, &xv));
    1217:   PetscCall(VecGetKokkosView(yin, &yv));
    1218:   DotNorm2 dn(xv, yv);
    1219:   PetscCallCXX(Kokkos::parallel_reduce(Kokkos::RangePolicy<>(PetscGetKokkosExecutionSpace(), 0, xin->map->n), dn, result));
    1220:   *dp = result[0];
    1221:   *nm = result[1];
    1222:   PetscCall(VecRestoreKokkosView(yin, &yv));
    1223:   PetscCall(VecRestoreKokkosView(xin, &xv));
    1224:   PetscCall(PetscLogGpuTimeEnd());
    1225:   PetscCall(PetscLogGpuFlops(4.0 * xin->map->n));
    1226:   PetscFunctionReturn(PETSC_SUCCESS);
    1227: }

    1229: PetscErrorCode VecConjugate_SeqKokkos(Vec xin)
    1230: {
    1231:   PetscScalarKokkosView xv;

    1233:   PetscFunctionBegin;
    1234:   PetscCall(PetscLogGpuTimeBegin());
    1235:   PetscCall(VecGetKokkosView(xin, &xv));
    1236:   PetscCallCXX(Kokkos::parallel_for(Kokkos::RangePolicy<>(PetscGetKokkosExecutionSpace(), 0, xin->map->n), KOKKOS_LAMBDA(const PetscInt &i) { xv(i) = PetscConj(xv(i)); }));
    1237:   PetscCall(VecRestoreKokkosView(xin, &xv));
    1238:   PetscCall(PetscLogGpuTimeEnd());
    1239:   PetscFunctionReturn(PETSC_SUCCESS);
    1240: }

    1242: /* Temporarily replace the array in vin with a[]. Return to the original array with a call to VecResetArray() */
    1243: PetscErrorCode VecPlaceArray_SeqKokkos(Vec vin, const PetscScalar *a)
    1244: {
    1245:   Vec_Seq    *vecseq = (Vec_Seq *)vin->data;
    1246:   Vec_Kokkos *veckok = static_cast<Vec_Kokkos *>(vin->spptr);

    1248:   PetscFunctionBegin;
    1249:   PetscCall(VecPlaceArray_Seq(vin, a));
    1250:   PetscCall(veckok->UpdateArray<HostMirrorMemorySpace>(vecseq->array));
    1251:   PetscFunctionReturn(PETSC_SUCCESS);
    1252: }

    1254: PetscErrorCode VecResetArray_SeqKokkos(Vec vin)
    1255: {
    1256:   Vec_Seq    *vecseq = (Vec_Seq *)vin->data;
    1257:   Vec_Kokkos *veckok = static_cast<Vec_Kokkos *>(vin->spptr);

    1259:   PetscFunctionBegin;
    1260:   /* User wants to unhook the provided host array. Sync it so that user can get the latest */
    1261:   PetscCall(KokkosDualViewSyncHost(veckok->v_dual, PetscGetKokkosExecutionSpace()));
    1262:   PetscCall(VecResetArray_Seq(vin)); /* Swap back the old host array, assuming its has the latest value */
    1263:   PetscCall(veckok->UpdateArray<HostMirrorMemorySpace>(vecseq->array));
    1264:   PetscFunctionReturn(PETSC_SUCCESS);
    1265: }

    1267: /*@C
    1268:   VecKokkosPlaceArray - Allows one to replace the device array in a VecKokkos vector with a
    1269:   device array provided by the user. This is useful to avoid copying an array into a vector.

    1271:   Logically Collective; No Fortran Support

    1273:   Input Parameters:
    1274: + v - the VecKokkos vector
    1275: - a - the device array

    1277:   Level: developer

    1279:   Notes:

    1281:   You can return to the original array with a call to `VecKokkosResetArray()`. `vec` does not take
    1282:   ownership of `array` in any way.

    1284:   The user manages the device array so PETSc doesn't care how it was allocated.

    1286:   The user must free `array` themselves but be careful not to
    1287:   do so before the vector has either been destroyed, had its original array restored with
    1288:   `VecKokkosResetArray()` or permanently replaced with `VecReplaceArray()`.

    1290: .seealso: [](ch_vectors), `Vec`, `VecGetArray()`, `VecKokkosResetArray()`, `VecReplaceArray()`, `VecResetArray()`
    1291: @*/
    1292: PetscErrorCode VecKokkosPlaceArray(Vec v, PetscScalar *a)
    1293: {
    1294:   Vec_Kokkos *veckok = static_cast<Vec_Kokkos *>(v->spptr);

    1296:   PetscFunctionBegin;
    1297:   VecErrorIfNotKokkos(v);
    1298:   // Sync the old device view before replacing it; so that when it is put back, it has the saved value.
    1299:   PetscCall(KokkosDualViewSyncDevice(veckok->v_dual, PetscGetKokkosExecutionSpace()));
    1300:   PetscCallCXX(veckok->unplaced_d = veckok->v_dual.view_device());
    1301:   // We assume a[] contains the latest data and discard the vector's old sync state
    1302:   PetscCall(veckok->UpdateArray<DefaultMemorySpace>(a));
    1303:   PetscFunctionReturn(PETSC_SUCCESS);
    1304: }

    1306: /*@C
    1307:   VecKokkosResetArray - Resets a vector to use its default memory. Call this
    1308:   after the use of `VecKokkosPlaceArray()`.

    1310:   Not Collective

    1312:   Input Parameter:
    1313: . v - the vector

    1315:   Level: developer

    1317:   Notes:

    1319:   After the call, the original array placed in with `VecKokkosPlaceArray()` will contain the latest value of the vector.
    1320:   Note that device kernels are asynchronous. Users are responsible to sync the device if they wish to have immediate access
    1321:   to the data in the array. Also, after the call, `v` will contain whatever data before `VecKokkosPlaceArray()`.

    1323: .seealso: [](ch_vectors), `Vec`, `VecGetArray()`, `VecRestoreArray()`, `VecReplaceArray()`, `VecKokkosPlaceArray()`
    1324: @*/
    1325: PetscErrorCode VecKokkosResetArray(Vec v)
    1326: {
    1327:   Vec_Kokkos *veckok = static_cast<Vec_Kokkos *>(v->spptr);

    1329:   PetscFunctionBegin;
    1330:   VecErrorIfNotKokkos(v);
    1331:   // User wants to unhook the provided device array. Sync it so that user can get the latest
    1332:   PetscCall(KokkosDualViewSyncDevice(veckok->v_dual, PetscGetKokkosExecutionSpace()));
    1333:   // Put the unplaced device array back, and set an appropriate modify flag
    1334:   PetscCall(veckok->UpdateArray<DefaultMemorySpace>(veckok->unplaced_d.data()));
    1335:   PetscFunctionReturn(PETSC_SUCCESS);
    1336: }

    1338: /* Replace the array in vin with a[] that must be allocated by PetscMalloc. a[] is owned by vin afterwards. */
    1339: PetscErrorCode VecReplaceArray_SeqKokkos(Vec vin, const PetscScalar *a)
    1340: {
    1341:   Vec_Seq    *vecseq = (Vec_Seq *)vin->data;
    1342:   Vec_Kokkos *veckok = static_cast<Vec_Kokkos *>(vin->spptr);

    1344:   PetscFunctionBegin;
    1345:   /* Make sure the users array has the latest values */
    1346:   if (vecseq->array != vecseq->array_allocated) PetscCall(KokkosDualViewSyncHost(veckok->v_dual, PetscGetKokkosExecutionSpace()));
    1347:   PetscCall(VecReplaceArray_Seq(vin, a));
    1348:   PetscCall(veckok->UpdateArray<HostMirrorMemorySpace>(vecseq->array));
    1349:   PetscFunctionReturn(PETSC_SUCCESS);
    1350: }

    1352: /* Maps the local portion of vector v into vector w */
    1353: PetscErrorCode VecGetLocalVector_SeqKokkos(Vec v, Vec w)
    1354: {
    1355:   PetscBool isKokkos = PETSC_FALSE;

    1357:   PetscFunctionBegin;
    1358:   PetscCall(PetscObjectTypeCompareAny((PetscObject)(w), &isKokkos, VECSEQKOKKOS, VECMPIKOKKOS, ""));
    1359:   if (!isKokkos) {
    1360:     PetscScalar *a;

    1362:     PetscCall(VecGetArray(v, &a));
    1363:     PetscCall(VecPlaceArray(w, a));
    1364:   } else { // The whole purpose of VecGetLocalVector is to quickly put v's data pointers into w (i.e., the v_dual, array fields)
    1365:     Vec_Kokkos *vkok = static_cast<Vec_Kokkos *>(v->spptr);
    1366:     Vec_Kokkos *wkok = static_cast<Vec_Kokkos *>(w->spptr);
    1367:     Vec_Seq    *vseq = static_cast<Vec_Seq *>(v->data);
    1368:     Vec_Seq    *wseq = static_cast<Vec_Seq *>(w->data);

    1370:     wkok->unplaced_dual = wkok->v_dual;
    1371:     wkok->v_dual        = vkok->v_dual;
    1372:     wseq->array         = vseq->array; // Need to also update w's host array pointer
    1373:   }
    1374:   PetscFunctionReturn(PETSC_SUCCESS);
    1375: }

    1377: PetscErrorCode VecRestoreLocalVector_SeqKokkos(Vec v, Vec w)
    1378: {
    1379:   PetscBool isKokkos = PETSC_FALSE;

    1381:   PetscFunctionBegin;
    1382:   PetscCall(PetscObjectTypeCompareAny((PetscObject)(w), &isKokkos, VECSEQKOKKOS, VECMPIKOKKOS, ""));
    1383:   if (!isKokkos) {
    1384:     PetscScalar *a;

    1386:     PetscCall(VecGetArray(w, &a));
    1387:     PetscCall(VecRestoreArray(v, &a));
    1388:     PetscCall(VecResetArray(w));
    1389:   } else {
    1390:     Vec_Kokkos *vkok = static_cast<Vec_Kokkos *>(v->spptr);
    1391:     Vec_Kokkos *wkok = static_cast<Vec_Kokkos *>(w->spptr);
    1392:     Vec_Seq    *wseq = static_cast<Vec_Seq *>(w->data);

    1394:     vkok->v_dual = wkok->v_dual;
    1395:     wkok->v_dual = wkok->unplaced_dual;
    1396:     wseq->array  = wkok->v_dual.view_host().data(); // restore w's host array pointer
    1397:   }
    1398:   PetscFunctionReturn(PETSC_SUCCESS);
    1399: }

    1401: /* Maps the local portion of vector v into vector w */
    1402: PetscErrorCode VecGetLocalVectorRead_SeqKokkos(Vec v, Vec w)
    1403: {
    1404:   PetscBool isKokkos = PETSC_FALSE;

    1406:   PetscFunctionBegin;
    1407:   PetscCall(PetscObjectTypeCompareAny((PetscObject)(w), &isKokkos, VECSEQKOKKOS, VECMPIKOKKOS, ""));
    1408:   if (!isKokkos) {
    1409:     const PetscScalar *a;

    1411:     PetscCall(VecGetArrayRead(v, &a));
    1412:     PetscCall(VecPlaceArray(w, a));
    1413:   } else {
    1414:     Vec_Kokkos *vkok = static_cast<Vec_Kokkos *>(v->spptr);
    1415:     Vec_Kokkos *wkok = static_cast<Vec_Kokkos *>(w->spptr);
    1416:     Vec_Seq    *vseq = static_cast<Vec_Seq *>(v->data);
    1417:     Vec_Seq    *wseq = static_cast<Vec_Seq *>(w->data);

    1419:     wkok->unplaced_dual = wkok->v_dual;
    1420:     wkok->v_dual        = vkok->v_dual;
    1421:     wseq->array         = vseq->array;
    1422:   }
    1423:   PetscFunctionReturn(PETSC_SUCCESS);
    1424: }

    1426: PetscErrorCode VecRestoreLocalVectorRead_SeqKokkos(Vec v, Vec w)
    1427: {
    1428:   PetscBool isKokkos = PETSC_FALSE;

    1430:   PetscFunctionBegin;
    1431:   PetscCall(PetscObjectTypeCompareAny((PetscObject)(w), &isKokkos, VECSEQKOKKOS, VECMPIKOKKOS, ""));
    1432:   if (!isKokkos) {
    1433:     const PetscScalar *a;

    1435:     PetscCall(VecGetArrayRead(w, &a));
    1436:     PetscCall(VecRestoreArrayRead(v, &a));
    1437:     PetscCall(VecResetArray(w));
    1438:   } else {
    1439:     Vec_Kokkos *vkok = static_cast<Vec_Kokkos *>(v->spptr);
    1440:     Vec_Kokkos *wkok = static_cast<Vec_Kokkos *>(w->spptr);
    1441:     Vec_Seq    *wseq = static_cast<Vec_Seq *>(w->data);

    1443:     vkok->v_dual = wkok->v_dual;
    1444:     wkok->v_dual = wkok->unplaced_dual;
    1445:     wseq->array  = wkok->v_dual.view_host().data();
    1446:   }
    1447:   PetscFunctionReturn(PETSC_SUCCESS);
    1448: }

    1450: PetscErrorCode VecGetArray_SeqKokkos(Vec v, PetscScalar **a)
    1451: {
    1452:   Vec_Kokkos *veckok = static_cast<Vec_Kokkos *>(v->spptr);

    1454:   PetscFunctionBegin;
    1455:   PetscCall(KokkosDualViewSyncHost(veckok->v_dual, PetscGetKokkosExecutionSpace()));
    1456:   *a = *((PetscScalar **)v->data);
    1457:   PetscFunctionReturn(PETSC_SUCCESS);
    1458: }

    1460: PetscErrorCode VecRestoreArray_SeqKokkos(Vec v, PetscScalar **)
    1461: {
    1462:   Vec_Kokkos *veckok = static_cast<Vec_Kokkos *>(v->spptr);

    1464:   PetscFunctionBegin;
    1465:   PetscCallCXX(veckok->v_dual.modify_host());
    1466:   PetscFunctionReturn(PETSC_SUCCESS);
    1467: }

    1469: /* Get array on host to overwrite, so no need to sync host. In VecRestoreArrayWrite() we will mark host is modified. */
    1470: PetscErrorCode VecGetArrayWrite_SeqKokkos(Vec v, PetscScalar **a)
    1471: {
    1472:   Vec_Kokkos *veckok = static_cast<Vec_Kokkos *>(v->spptr);

    1474:   PetscFunctionBegin;
    1475:   PetscCallCXX(veckok->v_dual.clear_sync_state());
    1476:   *a = veckok->v_dual.view_host().data();
    1477:   PetscFunctionReturn(PETSC_SUCCESS);
    1478: }

    1480: PetscErrorCode VecGetArrayAndMemType_SeqKokkos(Vec v, PetscScalar **a, PetscMemType *mtype)
    1481: {
    1482:   Vec_Kokkos *veckok = static_cast<Vec_Kokkos *>(v->spptr);

    1484:   PetscFunctionBegin;
    1485:   /* Always return up-to-date in the default memory space */
    1486:   PetscCall(KokkosDualViewSyncDevice(veckok->v_dual, PetscGetKokkosExecutionSpace()));
    1487:   *a = veckok->v_dual.view_device().data();
    1488:   if (mtype) *mtype = PETSC_MEMTYPE_KOKKOS; // Could be PETSC_MEMTYPE_HOST when Kokkos was not configured with cuda etc.
    1489:   PetscFunctionReturn(PETSC_SUCCESS);
    1490: }

    1492: PetscErrorCode VecRestoreArrayAndMemType_SeqKokkos(Vec v, PetscScalar **)
    1493: {
    1494:   Vec_Kokkos *veckok = static_cast<Vec_Kokkos *>(v->spptr);

    1496:   PetscFunctionBegin;
    1497:   if (PetscMemTypeHost(PETSC_MEMTYPE_KOKKOS)) PetscCallCXX(veckok->v_dual.modify_host());
    1498:   else PetscCallCXX(veckok->v_dual.modify_device());
    1499:   PetscFunctionReturn(PETSC_SUCCESS);
    1500: }

    1502: PetscErrorCode VecGetArrayWriteAndMemType_SeqKokkos(Vec v, PetscScalar **a, PetscMemType *mtype)
    1503: {
    1504:   Vec_Kokkos *veckok = static_cast<Vec_Kokkos *>(v->spptr);

    1506:   PetscFunctionBegin;
    1507:   // Since the data will be overwritten, we clear the sync state to suppress potential memory copying in sync'ing
    1508:   PetscCallCXX(veckok->v_dual.clear_sync_state()); // So that in restore, we can safely modify_device()
    1509:   PetscCall(KokkosDualViewSyncDevice(veckok->v_dual, PetscGetKokkosExecutionSpace()));
    1510:   *a = veckok->v_dual.view_device().data();
    1511:   if (mtype) *mtype = PETSC_MEMTYPE_KOKKOS; // Could be PETSC_MEMTYPE_HOST when Kokkos was not configured with cuda etc.
    1512:   PetscFunctionReturn(PETSC_SUCCESS);
    1513: }

    1515: /* Copy xin's sync state to y */
    1516: static PetscErrorCode VecCopySyncState_Kokkos_Private(Vec xin, Vec yout)
    1517: {
    1518:   Vec_Kokkos *xkok = static_cast<Vec_Kokkos *>(xin->spptr);
    1519:   Vec_Kokkos *ykok = static_cast<Vec_Kokkos *>(yout->spptr);

    1521:   PetscFunctionBegin;
    1522:   PetscCallCXX(ykok->v_dual.clear_sync_state());
    1523:   if (xkok->v_dual.need_sync_host()) PetscCallCXX(ykok->v_dual.modify_device());
    1524:   else if (xkok->v_dual.need_sync_device()) PetscCallCXX(ykok->v_dual.modify_host());
    1525:   PetscFunctionReturn(PETSC_SUCCESS);
    1526: }

    1528: static PetscErrorCode VecCreateSeqKokkosWithArrays_Private(MPI_Comm, PetscInt, PetscInt, const PetscScalar[], const PetscScalar[], Vec *);

    1530: /* Internal routine shared by VecGetSubVector_{SeqKokkos,MPIKokkos} */
    1531: PetscErrorCode VecGetSubVector_Kokkos_Private(Vec x, PetscBool xIsMPI, IS is, Vec *y)
    1532: {
    1533:   PetscBool contig;
    1534:   PetscInt  n, N, start, bs;
    1535:   MPI_Comm  comm;
    1536:   Vec       z;

    1538:   PetscFunctionBegin;
    1539:   PetscCall(PetscObjectGetComm((PetscObject)x, &comm));
    1540:   PetscCall(ISGetLocalSize(is, &n));
    1541:   PetscCall(ISGetSize(is, &N));
    1542:   PetscCall(VecGetSubVectorContiguityAndBS_Private(x, is, &contig, &start, &bs));

    1544:   if (contig) { /* We can do a no-copy (in-place) implementation with y sharing x's arrays */
    1545:     Vec_Kokkos        *xkok    = static_cast<Vec_Kokkos *>(x->spptr);
    1546:     const PetscScalar *array_h = xkok->v_dual.view_host().data() + start;
    1547:     const PetscScalar *array_d = xkok->v_dual.view_device().data() + start;

    1549:     /* These calls assume the input arrays are synced */
    1550:     if (xIsMPI) PetscCall(VecCreateMPIKokkosWithArrays_Private(comm, bs, n, N, array_h, array_d, &z)); /* x could be MPI even when x's comm size = 1 */
    1551:     else PetscCall(VecCreateSeqKokkosWithArrays_Private(comm, bs, n, array_h, array_d, &z));

    1553:     PetscCall(VecCopySyncState_Kokkos_Private(x, z)); /* Copy x's sync state to z */

    1555:     /* This is relevant only in debug mode */
    1556:     PetscInt state = 0;
    1557:     PetscCall(VecLockGet(x, &state));
    1558:     if (state) { /* x is either in read or read/write mode, therefore z, overlapped with x, can only be in read mode */
    1559:       PetscCall(VecLockReadPush(z));
    1560:     }

    1562:     z->ops->placearray   = NULL; /* z's arrays can't be replaced, because z does not own them */
    1563:     z->ops->replacearray = NULL;

    1565:   } else { /* Have to create a VecScatter and a stand-alone vector */
    1566:     PetscCall(VecGetSubVectorThroughVecScatter_Private(x, is, bs, &z));
    1567:   }
    1568:   *y = z;
    1569:   PetscFunctionReturn(PETSC_SUCCESS);
    1570: }

    1572: static PetscErrorCode VecGetSubVector_SeqKokkos(Vec x, IS is, Vec *y)
    1573: {
    1574:   PetscFunctionBegin;
    1575:   PetscCall(VecGetSubVector_Kokkos_Private(x, PETSC_FALSE, is, y));
    1576:   PetscFunctionReturn(PETSC_SUCCESS);
    1577: }

    1579: /* Restore subvector y to x */
    1580: PetscErrorCode VecRestoreSubVector_SeqKokkos(Vec x, IS, Vec *y)
    1581: {
    1582:   VecScatter                    vscat;
    1583:   PETSC_UNUSED PetscObjectState dummystate = 0;
    1584:   PetscBool                     unchanged;

    1586:   PetscFunctionBegin;
    1587:   PetscCall(PetscObjectComposedDataGetInt((PetscObject)*y, VecGetSubVectorSavedStateId, dummystate, unchanged));
    1588:   if (unchanged) PetscFunctionReturn(PETSC_SUCCESS); /* If y's state has not changed since VecGetSubVector(), we only need to destroy it */

    1590:   PetscCall(PetscObjectQuery((PetscObject)*y, "VecGetSubVector_Scatter", (PetscObject *)&vscat));
    1591:   if (vscat) {
    1592:     PetscCall(VecScatterBegin(vscat, *y, x, INSERT_VALUES, SCATTER_REVERSE));
    1593:     PetscCall(VecScatterEnd(vscat, *y, x, INSERT_VALUES, SCATTER_REVERSE));
    1594:   } else { /* y and x's (host and device) arrays overlap */
    1595:     Vec_Kokkos *xkok = static_cast<Vec_Kokkos *>(x->spptr);
    1596:     Vec_Kokkos *ykok = static_cast<Vec_Kokkos *>((*y)->spptr);
    1597:     PetscInt    state;

    1599:     PetscCall(VecLockGet(x, &state));
    1600:     PetscCheck(!state, PetscObjectComm((PetscObject)x), PETSC_ERR_ARG_WRONGSTATE, "Vec x is locked for read-only or read/write access");

    1602:     /* The tricky part: one has to carefully sync the arrays */
    1603:     auto exec = PetscGetKokkosExecutionSpace();
    1604:     if (xkok->v_dual.need_sync_device()) { /* x's host has newer data */
    1605:       /* Move y's latest values to host (since y is just a subset of x) */
    1606:       PetscCall(KokkosDualViewSyncHost(ykok->v_dual, exec));
    1607:     } else if (xkok->v_dual.need_sync_host()) {                /* x's device has newer data */
    1608:       PetscCall(KokkosDualViewSyncDevice(ykok->v_dual, exec)); /* Move y's latest data to device */
    1609:     } else {                                                   /* x's host and device data is already sync'ed; Copy y's sync state to x */
    1610:       PetscCall(VecCopySyncState_Kokkos_Private(*y, x));
    1611:     }
    1612:     PetscCall(PetscObjectStateIncrease((PetscObject)x)); /* Since x is updated */
    1613:   }
    1614:   PetscFunctionReturn(PETSC_SUCCESS);
    1615: }

    1617: static PetscErrorCode VecSetPreallocationCOO_SeqKokkos(Vec x, PetscCount ncoo, const PetscInt coo_i[])
    1618: {
    1619:   Vec_Seq    *vecseq = static_cast<Vec_Seq *>(x->data);
    1620:   Vec_Kokkos *veckok = static_cast<Vec_Kokkos *>(x->spptr);
    1621:   PetscInt    m;

    1623:   PetscFunctionBegin;
    1624:   PetscCall(VecSetPreallocationCOO_Seq(x, ncoo, coo_i));
    1625:   PetscCall(VecGetLocalSize(x, &m));
    1626:   PetscCall(veckok->SetUpCOO(vecseq, m));
    1627:   PetscFunctionReturn(PETSC_SUCCESS);
    1628: }

    1630: static PetscErrorCode VecSetValuesCOO_SeqKokkos(Vec x, const PetscScalar v[], InsertMode imode)
    1631: {
    1632:   Vec_Seq                    *vecseq = static_cast<Vec_Seq *>(x->data);
    1633:   Vec_Kokkos                 *veckok = static_cast<Vec_Kokkos *>(x->spptr);
    1634:   const PetscCountKokkosView &jmap1  = veckok->jmap1_d;
    1635:   const PetscCountKokkosView &perm1  = veckok->perm1_d;
    1636:   PetscScalarKokkosView       xv; /* View for vector x */
    1637:   ConstPetscScalarKokkosView  vv; /* View for array v[] */
    1638:   PetscInt                    m;
    1639:   PetscMemType                memtype;

    1641:   PetscFunctionBegin;
    1642:   PetscCall(VecGetLocalSize(x, &m));
    1643:   PetscCall(PetscGetMemType(v, &memtype));
    1644:   if (PetscMemTypeHost(memtype)) { /* If user gave v[] in host, we might need to copy it to device if any */
    1645:     PetscCallCXX(vv = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), ConstPetscScalarKokkosViewHost(v, vecseq->coo_n)));
    1646:   } else {
    1647:     PetscCallCXX(vv = ConstPetscScalarKokkosView(v, vecseq->coo_n)); /* Directly use v[]'s memory */
    1648:   }

    1650:   if (imode == INSERT_VALUES) PetscCall(VecGetKokkosViewWrite(x, &xv)); /* write vector */
    1651:   else PetscCall(VecGetKokkosView(x, &xv));                             /* read & write vector */

    1653:   PetscCallCXX(Kokkos::parallel_for(
    1654:     Kokkos::RangePolicy<>(PetscGetKokkosExecutionSpace(), 0, m), KOKKOS_LAMBDA(const PetscInt &i) {
    1655:       PetscScalar sum = 0.0;
    1656:       for (PetscCount k = jmap1(i); k < jmap1(i + 1); k++) sum += vv(perm1(k));
    1657:       xv(i) = (imode == INSERT_VALUES ? 0.0 : xv(i)) + sum;
    1658:     }));

    1660:   if (imode == INSERT_VALUES) PetscCall(VecRestoreKokkosViewWrite(x, &xv));
    1661:   else PetscCall(VecRestoreKokkosView(x, &xv));
    1662:   PetscFunctionReturn(PETSC_SUCCESS);
    1663: }

    1665: static PetscErrorCode VecCreate_SeqKokkos_Common(Vec); // forward declaration

    1667: /* Duplicate layout etc but not the values in the input vector */
    1668: static PetscErrorCode VecDuplicate_SeqKokkos(Vec win, Vec *v)
    1669: {
    1670:   PetscFunctionBegin;
    1671:   PetscCall(VecDuplicate_Seq(win, v)); /* It also dups ops of win */
    1672:   PetscCall(VecCreate_SeqKokkos_Common(*v));
    1673:   (*v)->ops[0] = win->ops[0]; // recover the ops[]. We need to always follow ops in win
    1674:   PetscFunctionReturn(PETSC_SUCCESS);
    1675: }

    1677: static PetscErrorCode VecDestroy_SeqKokkos(Vec v)
    1678: {
    1679:   Vec_Kokkos *veckok = static_cast<Vec_Kokkos *>(v->spptr);
    1680:   Vec_Seq    *vecseq = static_cast<Vec_Seq *>(v->data);

    1682:   PetscFunctionBegin;
    1683:   delete veckok;
    1684:   v->spptr = NULL;
    1685:   if (vecseq) PetscCall(VecDestroy_Seq(v));
    1686:   PetscFunctionReturn(PETSC_SUCCESS);
    1687: }

    1689: PETSC_INTERN PetscErrorCode VecBindToCPU_SeqKokkos(Vec v, PetscBool flg)
    1690: {
    1691:   PetscFunctionBegin;
    1692:   PetscCheck(v->boundtocpu == flg, PetscObjectComm((PetscObject)v), PETSC_ERR_SUP, "Changing binding of a VECKOKKOS vector is not supported yet");
    1693:   PetscFunctionReturn(PETSC_SUCCESS);
    1694: }

    1696: // Shared by all VecCreate/Duplicate routines for VecSeqKokkos
    1697: static PetscErrorCode VecCreate_SeqKokkos_Common(Vec v)
    1698: {
    1699:   PetscFunctionBegin;
    1700:   v->boundtocpu           = PetscDefined(HAVE_KOKKOS_WITHOUT_GPU) ? PETSC_TRUE : PETSC_FALSE; // VECKOKKOS has yet to support CPU binding. But in this case, we deem it is bound to CPU.
    1701:   v->ops->bindtocpu       = VecBindToCPU_SeqKokkos;
    1702:   v->ops->abs             = VecAbs_SeqKokkos;
    1703:   v->ops->reciprocal      = VecReciprocal_SeqKokkos;
    1704:   v->ops->pointwisemult   = VecPointwiseMult_SeqKokkos;
    1705:   v->ops->min             = VecMin_SeqKokkos;
    1706:   v->ops->max             = VecMax_SeqKokkos;
    1707:   v->ops->sum             = VecSum_SeqKokkos;
    1708:   v->ops->shift           = VecShift_SeqKokkos;
    1709:   v->ops->norm            = VecNorm_SeqKokkos;
    1710:   v->ops->scale           = VecScale_SeqKokkos;
    1711:   v->ops->copy            = VecCopy_SeqKokkos;
    1712:   v->ops->set             = VecSet_SeqKokkos;
    1713:   v->ops->swap            = VecSwap_SeqKokkos;
    1714:   v->ops->axpy            = VecAXPY_SeqKokkos;
    1715:   v->ops->axpby           = VecAXPBY_SeqKokkos;
    1716:   v->ops->axpbypcz        = VecAXPBYPCZ_SeqKokkos;
    1717:   v->ops->pointwisedivide = VecPointwiseDivide_SeqKokkos;
    1718:   v->ops->setrandom       = VecSetRandom_SeqKokkos;

    1720:   v->ops->dot   = VecDot_SeqKokkos;
    1721:   v->ops->tdot  = VecTDot_SeqKokkos;
    1722:   v->ops->mdot  = VecMDot_SeqKokkos;
    1723:   v->ops->mtdot = VecMTDot_SeqKokkos;

    1725:   v->ops->dot_local   = VecDot_SeqKokkos;
    1726:   v->ops->tdot_local  = VecTDot_SeqKokkos;
    1727:   v->ops->mdot_local  = VecMDot_SeqKokkos;
    1728:   v->ops->mtdot_local = VecMTDot_SeqKokkos;

    1730:   v->ops->norm_local             = VecNorm_SeqKokkos;
    1731:   v->ops->maxpy                  = VecMAXPY_SeqKokkos;
    1732:   v->ops->aypx                   = VecAYPX_SeqKokkos;
    1733:   v->ops->waxpy                  = VecWAXPY_SeqKokkos;
    1734:   v->ops->dotnorm2               = VecDotNorm2_SeqKokkos;
    1735:   v->ops->errorwnorm             = VecErrorWeightedNorms_SeqKokkos;
    1736:   v->ops->placearray             = VecPlaceArray_SeqKokkos;
    1737:   v->ops->replacearray           = VecReplaceArray_SeqKokkos;
    1738:   v->ops->resetarray             = VecResetArray_SeqKokkos;
    1739:   v->ops->destroy                = VecDestroy_SeqKokkos;
    1740:   v->ops->duplicate              = VecDuplicate_SeqKokkos;
    1741:   v->ops->conjugate              = VecConjugate_SeqKokkos;
    1742:   v->ops->getlocalvector         = VecGetLocalVector_SeqKokkos;
    1743:   v->ops->restorelocalvector     = VecRestoreLocalVector_SeqKokkos;
    1744:   v->ops->getlocalvectorread     = VecGetLocalVectorRead_SeqKokkos;
    1745:   v->ops->restorelocalvectorread = VecRestoreLocalVectorRead_SeqKokkos;
    1746:   v->ops->getarraywrite          = VecGetArrayWrite_SeqKokkos;
    1747:   v->ops->getarray               = VecGetArray_SeqKokkos;
    1748:   v->ops->restorearray           = VecRestoreArray_SeqKokkos;

    1750:   v->ops->getarrayandmemtype      = VecGetArrayAndMemType_SeqKokkos;
    1751:   v->ops->restorearrayandmemtype  = VecRestoreArrayAndMemType_SeqKokkos;
    1752:   v->ops->getarraywriteandmemtype = VecGetArrayWriteAndMemType_SeqKokkos;
    1753:   v->ops->getsubvector            = VecGetSubVector_SeqKokkos;
    1754:   v->ops->restoresubvector        = VecRestoreSubVector_SeqKokkos;

    1756:   v->ops->setpreallocationcoo = VecSetPreallocationCOO_SeqKokkos;
    1757:   v->ops->setvaluescoo        = VecSetValuesCOO_SeqKokkos;

    1759:   v->offloadmask = PETSC_OFFLOAD_KOKKOS; // Mark this is a VECKOKKOS; We use this flag for cheap VECKOKKOS test.
    1760:   PetscFunctionReturn(PETSC_SUCCESS);
    1761: }

    1763: /*@C
    1764:   VecCreateSeqKokkosWithArray - Creates a Kokkos sequential array-style vector,
    1765:   where the user provides the array space to store the vector values. The array
    1766:   provided must be a device array.

    1768:   Collective

    1770:   Input Parameters:
    1771: + comm   - the communicator, should be PETSC_COMM_SELF
    1772: . bs     - the block size
    1773: . n      - the vector length
    1774: - darray - device memory where the vector elements are to be stored.

    1776:   Output Parameter:
    1777: . v - the vector

    1779:   Notes:
    1780:   Use VecDuplicate() or VecDuplicateVecs() to form additional vectors of the
    1781:   same type as an existing vector.

    1783:   PETSc does NOT free the array when the vector is destroyed via VecDestroy().
    1784:   The user should not free the array until the vector is destroyed.

    1786:   Level: intermediate

    1788: .seealso: `VecCreateMPICUDAWithArray()`, `VecCreate()`, `VecDuplicate()`, `VecDuplicateVecs()`,
    1789:           `VecCreateGhost()`, `VecCreateSeq()`, `VecCreateSeqWithArray()`,
    1790:           `VecCreateMPIWithArray()`
    1791: @*/
    1792: PetscErrorCode VecCreateSeqKokkosWithArray(MPI_Comm comm, PetscInt bs, PetscInt n, const PetscScalar darray[], Vec *v)
    1793: {
    1794:   PetscMPIInt  size;
    1795:   Vec          w;
    1796:   Vec_Kokkos  *veckok = NULL;
    1797:   PetscScalar *harray;

    1799:   PetscFunctionBegin;
    1800:   PetscCallMPI(MPI_Comm_size(comm, &size));
    1801:   PetscCheck(size <= 1, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Cannot create VECSEQKOKKOS on more than one process");

    1803:   PetscCall(PetscKokkosInitializeCheck());
    1804:   PetscCall(VecCreate(comm, &w));
    1805:   PetscCall(VecSetSizes(w, n, n));
    1806:   PetscCall(VecSetBlockSize(w, bs));
    1807:   if (!darray) { /* Allocate memory ourself if user provided NULL */
    1808:     PetscCall(VecSetType(w, VECSEQKOKKOS));
    1809:   } else {
    1810:     /* Build a VECSEQ, get its harray, and then build Vec_Kokkos along with darray */
    1811:     if (std::is_same<DefaultMemorySpace, HostMirrorMemorySpace>::value) {
    1812:       harray = const_cast<PetscScalar *>(darray);
    1813:       PetscCall(VecCreate_Seq_Private(w, harray)); /* Build a sequential vector with harray */
    1814:     } else {
    1815:       PetscCall(VecSetType(w, VECSEQ));
    1816:       harray = static_cast<Vec_Seq *>(w->data)->array;
    1817:     }
    1818:     PetscCall(PetscObjectChangeTypeName((PetscObject)w, VECSEQKOKKOS)); /* Change it to Kokkos */
    1819:     PetscCall(VecCreate_SeqKokkos_Common(w));
    1820:     PetscCallCXX(veckok = new Vec_Kokkos{n, harray, const_cast<PetscScalar *>(darray)});
    1821:     PetscCallCXX(veckok->v_dual.modify_device()); /* Mark the device is modified */
    1822:     w->spptr = static_cast<void *>(veckok);
    1823:   }
    1824:   *v = w;
    1825:   PetscFunctionReturn(PETSC_SUCCESS);
    1826: }

    1828: // Convert VECSEQ to VECSEQKOKKOS. One could provide a non-NULL array_d[] for memory on device. If that is the case, we assume
    1829: // array_d[] is already synced with the host array. If array_d is NULL, then PETSc will allocate the device memory.
    1830: PetscErrorCode VecConvert_Seq_SeqKokkos_inplace(Vec v, PetscScalar *array_d)
    1831: {
    1832:   Vec_Seq *vecseq;

    1834:   PetscFunctionBegin;
    1835:   PetscCall(PetscKokkosInitializeCheck());
    1836:   PetscCall(PetscLayoutSetUp(v->map));
    1837:   PetscCall(PetscObjectChangeTypeName((PetscObject)v, VECSEQKOKKOS));
    1838:   PetscCall(VecCreate_SeqKokkos_Common(v));
    1839:   PetscCheck(!v->spptr, PETSC_COMM_SELF, PETSC_ERR_PLIB, "v->spptr not NULL");
    1840:   vecseq = static_cast<Vec_Seq *>(v->data);
    1841:   PetscCallCXX(v->spptr = new Vec_Kokkos(v->map->n, vecseq->array, array_d));
    1842:   PetscFunctionReturn(PETSC_SUCCESS);
    1843: }

    1845: // Create a VECSEQKOKKOS with layout and arrays
    1846: static PetscErrorCode VecCreateSeqKokkosWithLayoutAndArrays_Private(PetscLayout map, const PetscScalar harray[], const PetscScalar darray[], Vec *v)
    1847: {
    1848:   Vec w;

    1850:   PetscFunctionBegin;
    1851:   if (map->n > 0) PetscCheck(darray, map->comm, PETSC_ERR_ARG_WRONG, "darray cannot be NULL");
    1852: #if defined(KOKKOS_ENABLE_UNIFIED_MEMORY)
    1853:   PetscCheck(harray == darray, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "harray and darray must be the same");
    1854: #endif
    1855:   PetscCall(VecCreateSeqWithLayoutAndArray_Private(map, harray, &w));
    1856:   PetscCall(PetscObjectChangeTypeName((PetscObject)w, VECSEQKOKKOS)); // Change it to VECKOKKOS
    1857:   PetscCall(VecCreate_SeqKokkos_Common(w));
    1858:   PetscCallCXX(w->spptr = new Vec_Kokkos(map->n, const_cast<PetscScalar *>(harray), const_cast<PetscScalar *>(darray)));
    1859:   *v = w;
    1860:   PetscFunctionReturn(PETSC_SUCCESS);
    1861: }

    1863: /*
    1864:    VecCreateSeqKokkosWithArrays_Private - Creates a Kokkos sequential array-style vector
    1865:    with user-provided arrays on host and device.

    1867:    Collective

    1869:    Input Parameter:
    1870: +  comm - the communicator, should be PETSC_COMM_SELF
    1871: .  bs - the block size
    1872: .  n - the vector length
    1873: .  harray - host memory where the vector elements are to be stored.
    1874: -  darray - device memory where the vector elements are to be stored.

    1876:    Output Parameter:
    1877: .  v - the vector

    1879:    Notes:
    1880:    Unlike VecCreate{Seq,MPI}CUDAWithArrays(), this routine is private since we do not expect users to use it directly.

    1882:    If there is no device, then harray and darray must be the same.
    1883:    If n is not zero, then harray and darray must be allocated.
    1884:    After the call, the created vector is supposed to be in a synchronized state, i.e.,
    1885:    we suppose harray and darray have the same data.

    1887:    PETSc does NOT free the array when the vector is destroyed via VecDestroy().
    1888:    Caller should not free the array until the vector is destroyed.
    1889: */
    1890: static PetscErrorCode VecCreateSeqKokkosWithArrays_Private(MPI_Comm comm, PetscInt bs, PetscInt n, const PetscScalar harray[], const PetscScalar darray[], Vec *v)
    1891: {
    1892:   PetscMPIInt size;
    1893:   PetscLayout map;

    1895:   PetscFunctionBegin;
    1896:   PetscCall(PetscKokkosInitializeCheck());
    1897:   PetscCallMPI(MPI_Comm_size(comm, &size));
    1898:   PetscCheck(size == 1, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Cannot create VECSEQKOKKOS on more than one process");
    1899:   PetscCall(PetscLayoutCreateFromSizes(comm, n, n, bs, &map));
    1900:   PetscCall(VecCreateSeqKokkosWithLayoutAndArrays_Private(map, harray, darray, v));
    1901:   PetscCall(PetscLayoutDestroy(&map));
    1902:   PetscFunctionReturn(PETSC_SUCCESS);
    1903: }

    1905: /* TODO: ftn-auto generates veckok.kokkosf.c */
    1906: /*@C
    1907:   VecCreateSeqKokkos - Creates a standard, sequential array-style vector.

    1909:   Collective

    1911:   Input Parameters:
    1912: + comm - the communicator, should be `PETSC_COMM_SELF`
    1913: - n    - the vector length

    1915:   Output Parameter:
    1916: . v - the vector

    1918:   Notes:
    1919:   Use VecDuplicate() or VecDuplicateVecs() to form additional vectors of the
    1920:   same type as an existing vector.

    1922:   Level: intermediate

    1924: .seealso: `VecCreateMPI()`, `VecCreate()`, `VecDuplicate()`, `VecDuplicateVecs()`, `VecCreateGhost()`
    1925:  @*/
    1926: PetscErrorCode VecCreateSeqKokkos(MPI_Comm comm, PetscInt n, Vec *v)
    1927: {
    1928:   PetscFunctionBegin;
    1929:   PetscCall(PetscKokkosInitializeCheck());
    1930:   PetscCall(VecCreate(comm, v));
    1931:   PetscCall(VecSetSizes(*v, n, n));
    1932:   PetscCall(VecSetType(*v, VECSEQKOKKOS)); /* Calls VecCreate_SeqKokkos */
    1933:   PetscFunctionReturn(PETSC_SUCCESS);
    1934: }

    1936: // Duplicate a VECSEQKOKKOS
    1937: static PetscErrorCode VecDuplicateVecs_SeqKokkos_GEMV(Vec w, PetscInt m, Vec *V[])
    1938: {
    1939:   PetscInt64                lda; // use 64-bit as we will do "m * lda"
    1940:   PetscScalar              *array_h, *array_d;
    1941:   PetscLayout               map;
    1942:   PetscScalarKokkosDualView w_dual;

    1944:   PetscFunctionBegin;
    1945:   PetscCall(PetscKokkosInitializeCheck()); // as we'll call kokkos_malloc()
    1946:   PetscCall(PetscMalloc1(m, V));
    1947:   PetscCall(VecGetLayout(w, &map));
    1948:   VecGetLocalSizeAligned(w, 64, &lda); // get in lda the 64-bytes aligned local size
    1949:   // See comments in VecCreate_SeqKokkos() on why we use DualView to allocate the memory
    1950:   PetscCallCXX(w_dual = PetscScalarKokkosDualView("VecDuplicateVecs", m * lda)); // Kokkos init's w_dual to zero

    1952:   // create the m vectors with raw arrays from v_dual
    1953:   array_h = w_dual.view_host().data();
    1954:   array_d = w_dual.view_device().data();
    1955:   for (PetscInt i = 0; i < m; i++) {
    1956:     Vec v;
    1957:     PetscCall(VecCreateSeqKokkosWithLayoutAndArrays_Private(map, &array_h[i * lda], &array_d[i * lda], &v));
    1958:     PetscCall(PetscObjectListDuplicate(((PetscObject)w)->olist, &((PetscObject)v)->olist));
    1959:     PetscCall(PetscFunctionListDuplicate(((PetscObject)w)->qlist, &((PetscObject)v)->qlist));
    1960:     v->ops[0] = w->ops[0];
    1961:     (*V)[i]   = v;
    1962:   }

    1964:   // let the first vector own the long DualView, so when it is destroyed it will free the v_dual
    1965:   if (m) {
    1966:     Vec v = (*V)[0];

    1968:     static_cast<Vec_Kokkos *>(v->spptr)->w_dual = w_dual; // stash the memory
    1969:     // disable replacearray of the first vector, as freeing its memory also frees others in the group.
    1970:     // But replacearray of others is ok, as they don't own their array.
    1971:     if (m > 1) v->ops->replacearray = VecReplaceArray_Default_GEMV_Error;
    1972:   }
    1973:   PetscFunctionReturn(PETSC_SUCCESS);
    1974: }

    1976: /*MC
    1977:    VECSEQKOKKOS - VECSEQKOKKOS = "seqkokkos" - The basic sequential vector, modified to use Kokkos

    1979:    Options Database Keys:
    1980: . -vec_type seqkokkos - sets the vector type to VECSEQKOKKOS during a call to VecSetFromOptions()

    1982:   Level: beginner

    1984: .seealso: `VecCreate()`, `VecSetType()`, `VecSetFromOptions()`, `VecCreateMPIWithArray()`, `VECMPI`, `VecType`, `VecCreateMPI()`
    1985: M*/
    1986: PetscErrorCode VecCreate_SeqKokkos(Vec v)
    1987: {
    1988:   PetscBool                 mdot_use_gemv  = PETSC_TRUE;
    1989:   PetscBool                 maxpy_use_gemv = PETSC_FALSE; // default is false as we saw bad performance with vendors' GEMV with tall skinny matrices.
    1990:   PetscScalarKokkosDualView v_dual;

    1992:   PetscFunctionBegin;
    1993:   PetscCall(PetscKokkosInitializeCheck());
    1994:   PetscCall(PetscLayoutSetUp(v->map));

    1996:   // Use DualView to allocate both the host array and the device array.
    1997:   // DualView first allocates the device array and then mirrors it to host.
    1998:   // With unified memory (e.g., on AMD MI300A APU), the two arrays are the same, with the host array
    1999:   // sharing the device array allocated by hipMalloc(), but not the other way around if we call
    2000:   // VecCreate_Seq() first and let the device array share the host array allocated by malloc().
    2001:   // hipMalloc() has an advantage over malloc() as it gives great binding and page size settings automatically, see
    2002:   // https://hpc.llnl.gov/documentation/user-guides/using-el-capitan-systems/introduction-and-quickstart/pro-tips
    2003:   PetscCallCXX(v_dual = PetscScalarKokkosDualView("v_dual", v->map->n)); // Kokkos init's v_dual to zero

    2005:   PetscCall(VecCreate_Seq_Private(v, v_dual.view_host().data()));
    2006:   PetscCall(PetscObjectChangeTypeName((PetscObject)v, VECSEQKOKKOS));
    2007:   PetscCall(VecCreate_SeqKokkos_Common(v));
    2008:   PetscCheck(!v->spptr, PETSC_COMM_SELF, PETSC_ERR_PLIB, "v->spptr not NULL");
    2009:   PetscCallCXX(v->spptr = new Vec_Kokkos(v_dual));

    2011:   PetscCall(PetscOptionsGetBool(NULL, NULL, "-vec_mdot_use_gemv", &mdot_use_gemv, NULL));
    2012:   PetscCall(PetscOptionsGetBool(NULL, NULL, "-vec_maxpy_use_gemv", &maxpy_use_gemv, NULL));

    2014:   // allocate multiple vectors together
    2015:   if (mdot_use_gemv || maxpy_use_gemv) v->ops[0].duplicatevecs = VecDuplicateVecs_SeqKokkos_GEMV;

    2017:   if (mdot_use_gemv) {
    2018:     v->ops[0].mdot        = VecMDot_SeqKokkos_GEMV;
    2019:     v->ops[0].mtdot       = VecMTDot_SeqKokkos_GEMV;
    2020:     v->ops[0].mdot_local  = VecMDot_SeqKokkos_GEMV;
    2021:     v->ops[0].mtdot_local = VecMTDot_SeqKokkos_GEMV;
    2022:   }
    2023:   if (maxpy_use_gemv) v->ops[0].maxpy = VecMAXPY_SeqKokkos_GEMV;
    2024:   PetscFunctionReturn(PETSC_SUCCESS);
    2025: }