Actual source code: veckok.kokkos.cxx

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

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

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

 16: template <class MemorySpace>
 17: PetscErrorCode VecGetKokkosView_Private(Vec v, PetscScalarKokkosViewType<MemorySpace> *kv, PetscBool overwrite)
 18: {
 19:   Vec_Kokkos *veckok = static_cast<Vec_Kokkos *>(v->spptr);

 21:   VecErrorIfNotKokkos(v);
 22:   if (!overwrite) veckok->v_dual.sync<MemorySpace>(); /* If overwrite=true, no need to sync the space, since caller will overwrite the data */
 23:   *kv = veckok->v_dual.view<MemorySpace>();
 24:   return 0;
 25: }

 27: template <class MemorySpace>
 28: PetscErrorCode VecRestoreKokkosView_Private(Vec v, PetscScalarKokkosViewType<MemorySpace> *kv, PetscBool overwrite)
 29: {
 30:   Vec_Kokkos *veckok = static_cast<Vec_Kokkos *>(v->spptr);

 32:   VecErrorIfNotKokkos(v);
 33:   if (overwrite) veckok->v_dual.clear_sync_state(); /* If overwrite=true, clear the old sync state since user forced an overwrite */
 34:   veckok->v_dual.modify<MemorySpace>();
 35:   PetscObjectStateIncrease((PetscObject)v);
 36:   return 0;
 37: }

 39: template <class MemorySpace>
 40: PetscErrorCode VecGetKokkosView(Vec v, ConstPetscScalarKokkosViewType<MemorySpace> *kv)
 41: {
 42:   Vec_Kokkos *veckok = static_cast<Vec_Kokkos *>(v->spptr);
 43:   VecErrorIfNotKokkos(v);
 44:   veckok->v_dual.sync<MemorySpace>(); /* Sync the space for caller to read */
 45:   *kv = veckok->v_dual.view<MemorySpace>();
 46:   return 0;
 47: }

 49: /* Function template explicit instantiation */
 50: template PETSC_VISIBILITY_PUBLIC PetscErrorCode VecGetKokkosView(Vec, ConstPetscScalarKokkosView *);
 51: template <>
 52: PETSC_VISIBILITY_PUBLIC PetscErrorCode VecGetKokkosView(Vec v, PetscScalarKokkosView *kv)
 53: {
 54:   return VecGetKokkosView_Private(v, kv, PETSC_FALSE);
 55: }
 56: template <>
 57: PETSC_VISIBILITY_PUBLIC PetscErrorCode VecRestoreKokkosView(Vec v, PetscScalarKokkosView *kv)
 58: {
 59:   return VecRestoreKokkosView_Private(v, kv, PETSC_FALSE);
 60: }
 61: template <>
 62: PETSC_VISIBILITY_PUBLIC PetscErrorCode VecGetKokkosViewWrite(Vec v, PetscScalarKokkosView *kv)
 63: {
 64:   return VecGetKokkosView_Private(v, kv, PETSC_TRUE);
 65: }
 66: template <>
 67: PETSC_VISIBILITY_PUBLIC PetscErrorCode VecRestoreKokkosViewWrite(Vec v, PetscScalarKokkosView *kv)
 68: {
 69:   return VecRestoreKokkosView_Private(v, kv, PETSC_TRUE);
 70: }

 72: #if !defined(KOKKOS_ENABLE_DEFAULT_DEVICE_TYPE_HOST) /* Get host views if the default memory space is not host space */
 73: template PETSC_VISIBILITY_PUBLIC PetscErrorCode VecGetKokkosView(Vec, ConstPetscScalarKokkosViewHost *);
 74: template <>
 75: PETSC_VISIBILITY_PUBLIC PetscErrorCode VecGetKokkosView(Vec v, PetscScalarKokkosViewHost *kv)
 76: {
 77:   return VecGetKokkosView_Private(v, kv, PETSC_FALSE);
 78: }
 79: template <>
 80: PETSC_VISIBILITY_PUBLIC PetscErrorCode VecRestoreKokkosView(Vec v, PetscScalarKokkosViewHost *kv)
 81: {
 82:   return VecRestoreKokkosView_Private(v, kv, PETSC_FALSE);
 83: }
 84: template <>
 85: PETSC_VISIBILITY_PUBLIC PetscErrorCode VecGetKokkosViewWrite(Vec v, PetscScalarKokkosViewHost *kv)
 86: {
 87:   return VecGetKokkosView_Private(v, kv, PETSC_TRUE);
 88: }
 89: template <>
 90: PETSC_VISIBILITY_PUBLIC PetscErrorCode VecRestoreKokkosViewWrite(Vec v, PetscScalarKokkosViewHost *kv)
 91: {
 92:   return VecRestoreKokkosView_Private(v, kv, PETSC_TRUE);
 93: }
 94: #endif

 96: PetscErrorCode VecSetRandom_SeqKokkos(Vec xin, PetscRandom r)
 97: {
 98:   const PetscInt n = xin->map->n;
 99:   PetscScalar   *xx;

101:   VecGetArrayWrite(xin, &xx); /* TODO: generate randoms directly on device */
102:   for (PetscInt i = 0; i < n; i++) PetscRandomGetValue(r, &xx[i]);
103:   VecRestoreArrayWrite(xin, &xx);
104:   return 0;
105: }

107: /* x = |x| */
108: PetscErrorCode VecAbs_SeqKokkos(Vec xin)
109: {
110:   PetscScalarKokkosView xv;

112:   PetscLogGpuTimeBegin();
113:   VecGetKokkosView(xin, &xv);
114:   KokkosBlas::abs(xv, xv);
115:   VecRestoreKokkosView(xin, &xv);
116:   PetscLogGpuTimeEnd();
117:   return 0;
118: }

120: /* x = 1/x */
121: PetscErrorCode VecReciprocal_SeqKokkos(Vec xin)
122: {
123:   PetscScalarKokkosView xv;

125:   PetscLogGpuTimeBegin();
126:   VecGetKokkosView(xin, &xv);
127:   Kokkos::parallel_for(
128:     xv.extent(0), KOKKOS_LAMBDA(const int64_t i) {
129:       if (xv(i) != (PetscScalar)0.0) xv(i) = (PetscScalar)1.0 / xv(i);
130:     });
131:   VecRestoreKokkosView(xin, &xv);
132:   PetscLogGpuTimeEnd();
133:   return 0;
134: }

136: PetscErrorCode VecMin_SeqKokkos(Vec xin, PetscInt *p, PetscReal *val)
137: {
138:   typedef Kokkos::MinLoc<PetscReal, PetscInt>::value_type MinLocValue_t;
139:   ConstPetscScalarKokkosView                              xv;
140:   MinLocValue_t                                           minloc;

142:   PetscLogGpuTimeBegin();
143:   VecGetKokkosView(xin, &xv);
144:   Kokkos::parallel_reduce(
145:     "VecMin", xin->map->n,
146:     KOKKOS_LAMBDA(PetscInt i, MinLocValue_t & lminloc) {
147:       if (PetscRealPart(xv(i)) < lminloc.val) {
148:         lminloc.val = PetscRealPart(xv(i));
149:         lminloc.loc = i;
150:       }
151:     },
152:     Kokkos::MinLoc<PetscReal, PetscInt>(minloc)); /* Kokkos will set minloc properly even if xin is zero-lengthed */
153:   if (p) *p = minloc.loc;
154:   *val = minloc.val;
155:   VecRestoreKokkosView(xin, &xv);
156:   PetscLogGpuTimeEnd();
157:   return 0;
158: }

160: PetscErrorCode VecMax_SeqKokkos(Vec xin, PetscInt *p, PetscReal *val)
161: {
162:   typedef Kokkos::MaxLoc<PetscReal, PetscInt>::value_type MaxLocValue_t;
163:   ConstPetscScalarKokkosView                              xv;
164:   MaxLocValue_t                                           maxloc;

166:   PetscLogGpuTimeBegin();
167:   VecGetKokkosView(xin, &xv);
168:   Kokkos::parallel_reduce(
169:     "VecMax", xin->map->n,
170:     KOKKOS_LAMBDA(PetscInt i, MaxLocValue_t & lmaxloc) {
171:       if (PetscRealPart(xv(i)) > lmaxloc.val) {
172:         lmaxloc.val = PetscRealPart(xv(i));
173:         lmaxloc.loc = i;
174:       }
175:     },
176:     Kokkos::MaxLoc<PetscReal, PetscInt>(maxloc));
177:   if (p) *p = maxloc.loc;
178:   *val = maxloc.val;
179:   VecRestoreKokkosView(xin, &xv);
180:   PetscLogGpuTimeEnd();
181:   return 0;
182: }

184: PetscErrorCode VecSum_SeqKokkos(Vec xin, PetscScalar *sum)
185: {
186:   ConstPetscScalarKokkosView xv;

188:   PetscLogGpuTimeBegin();
189:   VecGetKokkosView(xin, &xv);
190:   *sum = KokkosBlas::sum(xv);
191:   VecRestoreKokkosView(xin, &xv);
192:   PetscLogGpuTimeEnd();
193:   return 0;
194: }

196: PetscErrorCode VecShift_SeqKokkos(Vec xin, PetscScalar shift)
197: {
198:   PetscScalarKokkosView xv;

200:   PetscLogGpuTimeBegin();
201:   VecGetKokkosView(xin, &xv);
202:   Kokkos::parallel_for(
203:     "VecShift", xin->map->n, KOKKOS_LAMBDA(PetscInt i) { xv(i) += shift; });
204:   VecRestoreKokkosView(xin, &xv);
205:   PetscLogGpuTimeEnd();
206:   return 0;
207: }

209: /* y = alpha x + y */
210: PetscErrorCode VecAXPY_SeqKokkos(Vec yin, PetscScalar alpha, Vec xin)
211: {
212:   PetscBool                  xiskok, yiskok;
213:   PetscScalarKokkosView      yv;
214:   ConstPetscScalarKokkosView xv;

216:   if (alpha == (PetscScalar)0.0) return 0;
217:   if (yin == xin) {
218:     VecScale_SeqKokkos(yin, alpha + 1);
219:     return 0;
220:   }
221:   PetscObjectTypeCompareAny((PetscObject)xin, &xiskok, VECSEQKOKKOS, VECMPIKOKKOS, "");
222:   PetscObjectTypeCompareAny((PetscObject)yin, &yiskok, VECSEQKOKKOS, VECMPIKOKKOS, "");
223:   if (xiskok && yiskok) {
224:     PetscLogGpuTimeBegin();
225:     VecGetKokkosView(xin, &xv);
226:     VecGetKokkosView(yin, &yv);
227:     KokkosBlas::axpy(alpha, xv, yv);
228:     VecRestoreKokkosView(xin, &xv);
229:     VecRestoreKokkosView(yin, &yv);
230:     PetscLogGpuTimeEnd();
231:     PetscLogGpuFlops(2.0 * yin->map->n);
232:   } else {
233:     VecAXPY_Seq(yin, alpha, xin);
234:   }
235:   return 0;
236: }

238: /* y = x + beta y */
239: PetscErrorCode VecAYPX_SeqKokkos(Vec yin, PetscScalar beta, Vec xin)
240: {
241:   /* One needs to define KOKKOSBLAS_OPTIMIZATION_LEVEL_AXPBY > 2 to have optimizations for cases alpha/beta = 0,+/-1 */
242:   VecAXPBY_SeqKokkos(yin, 1.0, beta, xin);
243:   return 0;
244: }

246: /* z = y^T x */
247: PetscErrorCode VecTDot_SeqKokkos(Vec xin, Vec yin, PetscScalar *z)
248: {
249:   ConstPetscScalarKokkosView xv, yv;

251:   PetscLogGpuTimeBegin();
252:   VecGetKokkosView(xin, &xv);
253:   VecGetKokkosView(yin, &yv);
254:   Kokkos::parallel_reduce(
255:     "VecTDot", xin->map->n, KOKKOS_LAMBDA(int64_t i, PetscScalar & update) { update += yv(i) * xv(i); }, *z); /* Kokkos always overwrites z, so no need to init it */
256:   VecRestoreKokkosView(yin, &yv);
257:   VecRestoreKokkosView(xin, &xv);
258:   PetscLogGpuTimeEnd();
259:   if (xin->map->n > 0) PetscLogGpuFlops(2.0 * xin->map->n);
260:   return 0;
261: }

263: struct TransposeDotTag { };
264: struct ConjugateDotTag { };

266: struct MDotFunctor {
267:   /* Note the C++ notation for an array typedef */
268:   typedef PetscScalar                           value_type[];
269:   typedef ConstPetscScalarKokkosView::size_type size_type;

271:   /* Tell Kokkos the result array's number of entries. This must be a public value in the functor */
272:   const size_type            value_count;
273:   ConstPetscScalarKokkosView xv, yv[8];

275:   MDotFunctor(ConstPetscScalarKokkosView &xv, const PetscInt ny, /* Number of valid entries in yv[8]. 1 <= ny <= 8 */
276:               ConstPetscScalarKokkosView &yv0, ConstPetscScalarKokkosView &yv1, ConstPetscScalarKokkosView &yv2, ConstPetscScalarKokkosView &yv3, ConstPetscScalarKokkosView &yv4, ConstPetscScalarKokkosView &yv5, ConstPetscScalarKokkosView &yv6, ConstPetscScalarKokkosView &yv7) :
277:     value_count(ny), xv(xv)
278:   {
279:     yv[0] = yv0;
280:     yv[1] = yv1;
281:     yv[2] = yv2;
282:     yv[3] = yv3;
283:     yv[4] = yv4;
284:     yv[5] = yv5;
285:     yv[6] = yv6;
286:     yv[7] = yv7;
287:   }

289:   KOKKOS_INLINE_FUNCTION void operator()(TransposeDotTag, const size_type i, value_type sum) const
290:   {
291:     PetscScalar xval = xv(i);
292:     for (size_type j = 0; j < value_count; ++j) sum[j] += yv[j](i) * xval;
293:   }

295:   KOKKOS_INLINE_FUNCTION void operator()(ConjugateDotTag, const size_type i, value_type sum) const
296:   {
297:     PetscScalar xval = xv(i);
298:     for (size_type j = 0; j < value_count; ++j) sum[j] += PetscConj(yv[j](i)) * xval;
299:   }

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

307:   KOKKOS_INLINE_FUNCTION void init(TransposeDotTag, value_type sum) const { init(sum); }
308:   KOKKOS_INLINE_FUNCTION void init(ConjugateDotTag, value_type sum) const { init(sum); }

310:   KOKKOS_INLINE_FUNCTION void join(value_type dst, const value_type src) const
311:   {
312:     for (size_type j = 0; j < value_count; j++) dst[j] += src[j];
313:   }

315:   KOKKOS_INLINE_FUNCTION void init(value_type sum) const
316:   {
317:     for (size_type j = 0; j < value_count; j++) sum[j] = 0.0;
318:   }
319: };

321: template <class WorkTag>
322: PetscErrorCode VecMultiDot_Private(Vec xin, PetscInt nv, const Vec yin[], PetscScalar *z)
323: {
324:   PetscInt                   i, j, cur = 0, ngroup = nv / 8, rem = nv % 8, N = xin->map->n;
325:   ConstPetscScalarKokkosView xv, yv[8];
326:   PetscScalarKokkosViewHost  zv(z, nv);

328:   VecGetKokkosView(xin, &xv);
329:   for (i = 0; i < ngroup; i++) { /* 8 y's per group */
330:     for (j = 0; j < 8; j++) VecGetKokkosView(yin[cur + j], &yv[j]);
331:     MDotFunctor mdot(xv, 8, yv[0], yv[1], yv[2], yv[3], yv[4], yv[5], yv[6], yv[7]); /* Hope Kokkos make it asynchronous */
332:     Kokkos::parallel_reduce(Kokkos::RangePolicy<WorkTag>(0, N), mdot, Kokkos::subview(zv, Kokkos::pair<PetscInt, PetscInt>(cur, cur + 8)));
333:     for (j = 0; j < 8; j++) VecRestoreKokkosView(yin[cur + j], &yv[j]);
334:     cur += 8;
335:   }

337:   if (rem) { /* The remaining */
338:     for (j = 0; j < rem; j++) VecGetKokkosView(yin[cur + j], &yv[j]);
339:     MDotFunctor mdot(xv, rem, yv[0], yv[1], yv[2], yv[3], yv[4], yv[5], yv[6], yv[7]);
340:     Kokkos::parallel_reduce(Kokkos::RangePolicy<WorkTag>(0, N), mdot, Kokkos::subview(zv, Kokkos::pair<PetscInt, PetscInt>(cur, cur + rem)));
341:     for (j = 0; j < rem; j++) VecRestoreKokkosView(yin[cur + j], &yv[j]);
342:   }
343:   VecRestoreKokkosView(xin, &xv);
344:   Kokkos::fence(); /* If reduce is async, then we need this fence to make sure z is ready for use on host */
345:   return 0;
346: }

348: /* z[i] = (x,y_i) = y_i^H x */
349: PetscErrorCode VecMDot_SeqKokkos(Vec xin, PetscInt nv, const Vec yin[], PetscScalar *z)
350: {
351:   VecMultiDot_Private<ConjugateDotTag>(xin, nv, yin, z);
352:   return 0;
353: }

355: /* z[i] = (x,y_i) = y_i^T x */
356: PetscErrorCode VecMTDot_SeqKokkos(Vec xin, PetscInt nv, const Vec yin[], PetscScalar *z)
357: {
358:   VecMultiDot_Private<TransposeDotTag>(xin, nv, yin, z);
359:   return 0;
360: }

362: /* x[:] = alpha */
363: PetscErrorCode VecSet_SeqKokkos(Vec xin, PetscScalar alpha)
364: {
365:   PetscScalarKokkosView xv;

367:   PetscLogGpuTimeBegin();
368:   VecGetKokkosViewWrite(xin, &xv);
369:   KokkosBlas::fill(xv, alpha);
370:   VecRestoreKokkosViewWrite(xin, &xv);
371:   PetscLogGpuTimeEnd();
372:   return 0;
373: }

375: /* x = alpha x */
376: PetscErrorCode VecScale_SeqKokkos(Vec xin, PetscScalar alpha)
377: {
378:   PetscScalarKokkosView xv;

380:   if (alpha == (PetscScalar)0.0) {
381:     VecSet_SeqKokkos(xin, alpha);
382:   } else if (alpha != (PetscScalar)1.0) {
383:     PetscLogGpuTimeBegin();
384:     VecGetKokkosView(xin, &xv);
385:     KokkosBlas::scal(xv, alpha, xv);
386:     VecRestoreKokkosView(xin, &xv);
387:     PetscLogGpuTimeEnd();
388:     PetscLogGpuFlops(xin->map->n);
389:   }
390:   return 0;
391: }

393: /* z = y^H x */
394: PetscErrorCode VecDot_SeqKokkos(Vec xin, Vec yin, PetscScalar *z)
395: {
396:   ConstPetscScalarKokkosView xv, yv;

398:   PetscLogGpuTimeBegin();
399:   VecGetKokkosView(xin, &xv);
400:   VecGetKokkosView(yin, &yv);
401:   *z = KokkosBlas::dot(yv, xv); /* KokkosBlas::dot(a,b) takes conjugate of a */
402:   VecRestoreKokkosView(xin, &xv);
403:   VecRestoreKokkosView(yin, &yv);
404:   PetscLogGpuTimeEnd();
405:   if (xin->map->n > 0) PetscLogGpuFlops(2.0 * xin->map->n - 1);
406:   return 0;
407: }

409: /* y = x, where x is VECKOKKOS, but y may be not */
410: PetscErrorCode VecCopy_SeqKokkos(Vec xin, Vec yin)
411: {
412:   PetscLogGpuTimeBegin();
413:   if (xin != yin) {
414:     Vec_Kokkos *xkok = static_cast<Vec_Kokkos *>(xin->spptr);
415:     if (yin->offloadmask == PETSC_OFFLOAD_KOKKOS) {
416:       /* y is also a VecKokkos */
417:       Vec_Kokkos *ykok = static_cast<Vec_Kokkos *>(yin->spptr);
418:       /* Kokkos rule: if x's host has newer data, it will copy to y's host view; otherwise to y's device view
419:         In case x's host is newer, y's device is newer, it will error (though should not, I think). So we just
420:         clear y's sync state.
421:        */
422:       ykok->v_dual.clear_sync_state();
423:       Kokkos::deep_copy(ykok->v_dual, xkok->v_dual);
424:     } else {
425:       PetscScalar *yarray;
426:       VecGetArrayWrite(yin, &yarray);
427:       PetscScalarKokkosViewHost yv(yarray, yin->map->n);
428:       if (xkok->v_dual.need_sync_host()) { /* x's device has newer data */
429:         Kokkos::deep_copy(yv, xkok->v_dual.view_device());
430:       } else {
431:         Kokkos::deep_copy(yv, xkok->v_dual.view_host());
432:       }
433:       VecRestoreArrayWrite(yin, &yarray);
434:     }
435:   }
436:   PetscLogGpuTimeEnd();
437:   return 0;
438: }

440: /* y[i] <--> x[i] */
441: PetscErrorCode VecSwap_SeqKokkos(Vec xin, Vec yin)
442: {
443:   PetscScalarKokkosView xv, yv;

445:   if (xin != yin) {
446:     PetscLogGpuTimeBegin();
447:     VecGetKokkosView(xin, &xv);
448:     VecGetKokkosView(yin, &yv);
449:     Kokkos::parallel_for(
450:       xin->map->n, KOKKOS_LAMBDA(const int64_t i) {
451:         PetscScalar tmp = xv(i);
452:         xv(i)           = yv(i);
453:         yv(i)           = tmp;
454:       });
455:     VecRestoreKokkosView(xin, &xv);
456:     VecRestoreKokkosView(yin, &yv);
457:     PetscLogGpuTimeEnd();
458:   }
459:   return 0;
460: }

462: /*  w = alpha x + y */
463: PetscErrorCode VecWAXPY_SeqKokkos(Vec win, PetscScalar alpha, Vec xin, Vec yin)
464: {
465:   ConstPetscScalarKokkosView xv, yv;
466:   PetscScalarKokkosView      wv;

468:   if (alpha == (PetscScalar)0.0) {
469:     VecCopy_SeqKokkos(yin, win);
470:   } else {
471:     PetscLogGpuTimeBegin();
472:     VecGetKokkosViewWrite(win, &wv);
473:     VecGetKokkosView(xin, &xv);
474:     VecGetKokkosView(yin, &yv);
475:     Kokkos::parallel_for(
476:       win->map->n, KOKKOS_LAMBDA(const int64_t i) { wv(i) = alpha * xv(i) + yv(i); });
477:     VecRestoreKokkosView(xin, &xv);
478:     VecRestoreKokkosView(yin, &yv);
479:     VecRestoreKokkosViewWrite(win, &wv);
480:     PetscLogGpuTimeEnd();
481:     PetscLogGpuFlops(2 * win->map->n);
482:   }
483:   return 0;
484: }

486: struct MAXPYFunctor {
487:   typedef ConstPetscScalarKokkosView::size_type size_type;

489:   PetscScalarKokkosView      yv;
490:   PetscInt                   nx; /* Significent entries in a[8] and xv[8] */
491:   PetscScalar                a[8];
492:   ConstPetscScalarKokkosView xv[8];

494:   MAXPYFunctor(PetscScalarKokkosView yv, PetscInt nx, 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) :
495:     yv(yv), nx(nx)
496:   {
497:     a[0]  = a0;
498:     a[1]  = a1;
499:     a[2]  = a2;
500:     a[3]  = a3;
501:     a[4]  = a4;
502:     a[5]  = a5;
503:     a[6]  = a6;
504:     a[7]  = a7;
505:     xv[0] = xv0;
506:     xv[1] = xv1;
507:     xv[2] = xv2;
508:     xv[3] = xv3;
509:     xv[4] = xv4;
510:     xv[5] = xv5;
511:     xv[6] = xv6;
512:     xv[7] = xv7;
513:   }

515:   KOKKOS_INLINE_FUNCTION void operator()(const size_type i) const
516:   {
517:     for (PetscInt j = 0; j < nx; ++j) yv(i) += a[j] * xv[j](i);
518:   }
519: };

521: /*  y = y + sum alpha[i] x[i] */
522: PetscErrorCode VecMAXPY_SeqKokkos(Vec yin, PetscInt nv, const PetscScalar *alpha, Vec *xin)
523: {
524:   PetscInt                   i, j, cur = 0, ngroup = nv / 8, rem = nv % 8;
525:   PetscScalarKokkosView      yv;
526:   PetscScalar                a[8];
527:   ConstPetscScalarKokkosView xv[8];

529:   VecGetKokkosView(yin, &yv);
530:   for (i = 0; i < ngroup; i++) { /* 8 x's per group */
531:     for (j = 0; j < 8; j++) {    /* Fill the parameters */
532:       a[j] = alpha[cur + j];
533:       VecGetKokkosView(xin[cur + j], &xv[j]);
534:     }
535:     MAXPYFunctor maxpy(yv, 8, 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]);
536:     Kokkos::parallel_for(yin->map->n, maxpy);
537:     for (j = 0; j < 8; j++) VecRestoreKokkosView(xin[cur + j], &xv[j]);
538:     cur += 8;
539:   }

541:   if (rem) { /* The remaining */
542:     for (j = 0; j < rem; j++) {
543:       a[j] = alpha[cur + j];
544:       VecGetKokkosView(xin[cur + j], &xv[j]);
545:     }
546:     MAXPYFunctor maxpy(yv, rem, 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]);
547:     Kokkos::parallel_for(yin->map->n, maxpy);
548:     for (j = 0; j < rem; j++) VecRestoreKokkosView(xin[cur + j], &xv[j]);
549:   }
550:   VecRestoreKokkosView(yin, &yv);
551:   PetscLogGpuFlops(nv * 2.0 * yin->map->n);
552:   return 0;
553: }

555: /* y = alpha x + beta y */
556: PetscErrorCode VecAXPBY_SeqKokkos(Vec yin, PetscScalar alpha, PetscScalar beta, Vec xin)
557: {
558:   ConstPetscScalarKokkosView xv;
559:   PetscScalarKokkosView      yv;
560:   PetscBool                  xiskok, yiskok;

562:   PetscObjectTypeCompareAny((PetscObject)xin, &xiskok, VECSEQKOKKOS, VECMPIKOKKOS, "");
563:   PetscObjectTypeCompareAny((PetscObject)yin, &yiskok, VECSEQKOKKOS, VECMPIKOKKOS, "");
564:   if (xiskok && yiskok) {
565:     PetscLogGpuTimeBegin();
566:     VecGetKokkosView(xin, &xv);
567:     VecGetKokkosView(yin, &yv);
568:     KokkosBlas::axpby(alpha, xv, beta, yv);
569:     VecRestoreKokkosView(xin, &xv);
570:     VecRestoreKokkosView(yin, &yv);
571:     PetscLogGpuTimeEnd();
572:     if (alpha == (PetscScalar)0.0 || beta == (PetscScalar)0.0) {
573:       PetscLogGpuFlops(xin->map->n);
574:     } else if (beta == (PetscScalar)1.0 || alpha == (PetscScalar)1.0) {
575:       PetscLogGpuFlops(2.0 * xin->map->n);
576:     } else {
577:       PetscLogGpuFlops(3.0 * xin->map->n);
578:     }
579:   } else {
580:     VecAXPBY_Seq(yin, alpha, beta, xin);
581:   }
582:   return 0;
583: }

585: /* z = alpha x + beta y + gamma z */
586: PetscErrorCode VecAXPBYPCZ_SeqKokkos(Vec zin, PetscScalar alpha, PetscScalar beta, PetscScalar gamma, Vec xin, Vec yin)
587: {
588:   ConstPetscScalarKokkosView xv, yv;
589:   PetscScalarKokkosView      zv;

591:   PetscLogGpuTimeBegin();
592:   VecGetKokkosView(zin, &zv);
593:   VecGetKokkosView(xin, &xv);
594:   VecGetKokkosView(yin, &yv);
595:   KokkosBlas::update(alpha, xv, beta, yv, gamma, zv);
596:   VecRestoreKokkosView(xin, &xv);
597:   VecRestoreKokkosView(yin, &yv);
598:   VecRestoreKokkosView(zin, &zv);
599:   PetscLogGpuTimeEnd();
600:   PetscLogGpuFlops(zin->map->n * 5);
601:   return 0;
602: }

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

606:   w is of type VecKokkos, but x, y may be not.
607: */
608: PetscErrorCode VecPointwiseMult_SeqKokkos(Vec win, Vec xin, Vec yin)
609: {
610:   PetscInt n;

612:   PetscLogGpuTimeBegin();
613:   VecGetLocalSize(win, &n);
614:   if (xin->offloadmask != PETSC_OFFLOAD_KOKKOS || yin->offloadmask != PETSC_OFFLOAD_KOKKOS) {
615:     PetscScalarKokkosViewHost wv;
616:     const PetscScalar        *xp, *yp;
617:     VecGetArrayRead(xin, &xp);
618:     VecGetArrayRead(yin, &yp);
619:     VecGetKokkosViewWrite(win, &wv);

621:     ConstPetscScalarKokkosViewHost xv(xp, n), yv(yp, n);
622:     Kokkos::parallel_for(
623:       Kokkos::RangePolicy<Kokkos::DefaultHostExecutionSpace>(0, n), KOKKOS_LAMBDA(const PetscInt i) { wv(i) = xv(i) * yv(i); });

625:     VecRestoreArrayRead(xin, &xp);
626:     VecRestoreArrayRead(yin, &yp);
627:     VecRestoreKokkosViewWrite(win, &wv);
628:   } else {
629:     ConstPetscScalarKokkosView xv, yv;
630:     PetscScalarKokkosView      wv;

632:     VecGetKokkosViewWrite(win, &wv);
633:     VecGetKokkosView(xin, &xv);
634:     VecGetKokkosView(yin, &yv);
635:     Kokkos::parallel_for(
636:       n, KOKKOS_LAMBDA(const PetscInt i) { wv(i) = xv(i) * yv(i); });
637:     VecRestoreKokkosView(yin, &yv);
638:     VecRestoreKokkosView(xin, &xv);
639:     VecRestoreKokkosViewWrite(win, &wv);
640:   }
641:   PetscLogGpuTimeEnd();
642:   PetscLogGpuFlops(n);
643:   return 0;
644: }

646: /* w = x/y */
647: PetscErrorCode VecPointwiseDivide_SeqKokkos(Vec win, Vec xin, Vec yin)
648: {
649:   PetscInt n;

651:   PetscLogGpuTimeBegin();
652:   VecGetLocalSize(win, &n);
653:   if (xin->offloadmask != PETSC_OFFLOAD_KOKKOS || yin->offloadmask != PETSC_OFFLOAD_KOKKOS) {
654:     PetscScalarKokkosViewHost wv;
655:     const PetscScalar        *xp, *yp;
656:     VecGetArrayRead(xin, &xp);
657:     VecGetArrayRead(yin, &yp);
658:     VecGetKokkosViewWrite(win, &wv);

660:     ConstPetscScalarKokkosViewHost xv(xp, n), yv(yp, n);
661:     Kokkos::parallel_for(
662:       Kokkos::RangePolicy<Kokkos::DefaultHostExecutionSpace>(0, n), KOKKOS_LAMBDA(const PetscInt i) {
663:         if (yv(i) != 0.0) wv(i) = xv(i) / yv(i);
664:         else wv(i) = 0.0;
665:       });

667:     VecRestoreArrayRead(xin, &xp);
668:     VecRestoreArrayRead(yin, &yp);
669:     VecRestoreKokkosViewWrite(win, &wv);
670:   } else {
671:     ConstPetscScalarKokkosView xv, yv;
672:     PetscScalarKokkosView      wv;

674:     VecGetKokkosViewWrite(win, &wv);
675:     VecGetKokkosView(xin, &xv);
676:     VecGetKokkosView(yin, &yv);
677:     Kokkos::parallel_for(
678:       n, KOKKOS_LAMBDA(const PetscInt i) {
679:         if (yv(i) != 0.0) wv(i) = xv(i) / yv(i);
680:         else wv(i) = 0.0;
681:       });
682:     VecRestoreKokkosView(yin, &yv);
683:     VecRestoreKokkosView(xin, &xv);
684:     VecRestoreKokkosViewWrite(win, &wv);
685:   }
686:   PetscLogGpuTimeEnd();
687:   PetscLogGpuFlops(win->map->n);
688:   return 0;
689: }

691: PetscErrorCode VecNorm_SeqKokkos(Vec xin, NormType type, PetscReal *z)
692: {
693:   const PetscInt             n = xin->map->n;
694:   ConstPetscScalarKokkosView xv;

696:   if (type == NORM_1_AND_2) {
697:     VecNorm_SeqKokkos(xin, NORM_1, z);
698:     VecNorm_SeqKokkos(xin, NORM_2, z + 1);
699:   } else {
700:     PetscLogGpuTimeBegin();
701:     VecGetKokkosView(xin, &xv);
702:     if (type == NORM_2 || type == NORM_FROBENIUS) {
703:       *z = KokkosBlas::nrm2(xv);
704:       PetscLogGpuFlops(PetscMax(2.0 * n - 1, 0.0));
705:     } else if (type == NORM_1) {
706:       *z = KokkosBlas::nrm1(xv);
707:       PetscLogGpuFlops(PetscMax(n - 1.0, 0.0));
708:     } else if (type == NORM_INFINITY) {
709:       *z = KokkosBlas::nrminf(xv);
710:     }
711:     VecRestoreKokkosView(xin, &xv);
712:     PetscLogGpuTimeEnd();
713:   }
714:   return 0;
715: }

717: /* A functor for DotNorm2 so that we can compute dp and nm in one kernel */
718: struct DotNorm2 {
719:   typedef PetscScalar                           value_type[];
720:   typedef ConstPetscScalarKokkosView::size_type size_type;

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

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

727:   KOKKOS_INLINE_FUNCTION void operator()(const size_type i, value_type result) const
728:   {
729:     result[0] += PetscConj(yv_(i)) * xv_(i);
730:     result[1] += PetscConj(yv_(i)) * yv_(i);
731:   }

733:   KOKKOS_INLINE_FUNCTION void join(value_type dst, const value_type src) const
734:   {
735:     dst[0] += src[0];
736:     dst[1] += src[1];
737:   }

739:   KOKKOS_INLINE_FUNCTION void init(value_type result) const
740:   {
741:     result[0] = 0.0;
742:     result[1] = 0.0;
743:   }
744: };

746: /* dp = y^H x, nm = y^H y */
747: PetscErrorCode VecDotNorm2_SeqKokkos(Vec xin, Vec yin, PetscScalar *dp, PetscScalar *nm)
748: {
749:   ConstPetscScalarKokkosView xv, yv;
750:   PetscScalar                result[2];

752:   PetscLogGpuTimeBegin();
753:   VecGetKokkosView(xin, &xv);
754:   VecGetKokkosView(yin, &yv);
755:   DotNorm2 dn(xv, yv);
756:   Kokkos::parallel_reduce(xin->map->n, dn, result);
757:   *dp = result[0];
758:   *nm = result[1];
759:   VecRestoreKokkosView(yin, &yv);
760:   VecRestoreKokkosView(xin, &xv);
761:   PetscLogGpuTimeEnd();
762:   PetscLogGpuFlops(4.0 * xin->map->n);
763:   return 0;
764: }

766: PetscErrorCode VecConjugate_SeqKokkos(Vec xin)
767: {
768: #if defined(PETSC_USE_COMPLEX)
769:   PetscScalarKokkosView xv;

771:   PetscLogGpuTimeBegin();
772:   VecGetKokkosView(xin, &xv);
773:   Kokkos::parallel_for(
774:     xin->map->n, KOKKOS_LAMBDA(int64_t i) { xv(i) = PetscConj(xv(i)); });
775:   VecRestoreKokkosView(xin, &xv);
776:   PetscLogGpuTimeEnd();
777: #else
778: #endif
779:   return 0;
780: }

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

788:   VecPlaceArray_Seq(vin, a);
789:   veckok->UpdateArray<Kokkos::HostSpace>(vecseq->array);
790:   return 0;
791: }

793: PetscErrorCode VecResetArray_SeqKokkos(Vec vin)
794: {
795:   Vec_Seq    *vecseq = (Vec_Seq *)vin->data;
796:   Vec_Kokkos *veckok = static_cast<Vec_Kokkos *>(vin->spptr);

798:   veckok->v_dual.sync_host();        /* User wants to unhook the provided host array. Sync it so that user can get the latest */
799:   VecResetArray_Seq(vin); /* Swap back the old host array, assuming its has the latest value */
800:   veckok->UpdateArray<Kokkos::HostSpace>(vecseq->array);
801:   return 0;
802: }

804: /* Replace the array in vin with a[] that must be allocated by PetscMalloc. a[] is owned by vin afterwords. */
805: PetscErrorCode VecReplaceArray_SeqKokkos(Vec vin, const PetscScalar *a)
806: {
807:   Vec_Seq    *vecseq = (Vec_Seq *)vin->data;
808:   Vec_Kokkos *veckok = static_cast<Vec_Kokkos *>(vin->spptr);

810:   /* Make sure the users array has the latest values */
811:   if (vecseq->array != vecseq->array_allocated) veckok->v_dual.sync_host();
812:   VecReplaceArray_Seq(vin, a);
813:   veckok->UpdateArray<Kokkos::HostSpace>(vecseq->array);
814:   return 0;
815: }

817: /* Maps the local portion of vector v into vector w */
818: PetscErrorCode VecGetLocalVector_SeqKokkos(Vec v, Vec w)
819: {
820:   Vec_Seq    *vecseq = static_cast<Vec_Seq *>(w->data);
821:   Vec_Kokkos *veckok = static_cast<Vec_Kokkos *>(w->spptr);

824:   /* Destroy w->data, w->spptr */
825:   if (vecseq) {
826:     PetscFree(vecseq->array_allocated);
827:     PetscFree(w->data);
828:   }
829:   delete veckok;

831:   /* Replace with v's */
832:   w->data  = v->data;
833:   w->spptr = v->spptr;
834:   PetscObjectStateIncrease((PetscObject)w);
835:   return 0;
836: }

838: PetscErrorCode VecRestoreLocalVector_SeqKokkos(Vec v, Vec w)
839: {
841:   v->data  = w->data;
842:   v->spptr = w->spptr;
843:   PetscObjectStateIncrease((PetscObject)v);
844:   /* TODO: need to think if setting w->data/spptr to NULL is safe */
845:   w->data  = NULL;
846:   w->spptr = NULL;
847:   return 0;
848: }

850: PetscErrorCode VecGetArray_SeqKokkos(Vec v, PetscScalar **a)
851: {
852:   Vec_Kokkos *veckok = static_cast<Vec_Kokkos *>(v->spptr);

854:   veckok->v_dual.sync_host();
855:   *a = *((PetscScalar **)v->data);
856:   return 0;
857: }

859: PetscErrorCode VecRestoreArray_SeqKokkos(Vec v, PetscScalar **a)
860: {
861:   Vec_Kokkos *veckok = static_cast<Vec_Kokkos *>(v->spptr);

863:   veckok->v_dual.modify_host();
864:   return 0;
865: }

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

872:   veckok->v_dual.clear_sync_state();
873:   *a = veckok->v_dual.view_host().data();
874:   return 0;
875: }

877: PetscErrorCode VecGetArrayAndMemType_SeqKokkos(Vec v, PetscScalar **a, PetscMemType *mtype)
878: {
879:   Vec_Kokkos *veckok = static_cast<Vec_Kokkos *>(v->spptr);

881:   if (std::is_same<DefaultMemorySpace, Kokkos::HostSpace>::value) {
882:     *a = veckok->v_dual.view_host().data();
883:     if (mtype) *mtype = PETSC_MEMTYPE_HOST;
884:   } else {
885:     /* When there is device, we always return up-to-date device data */
886:     veckok->v_dual.sync_device();
887:     *a = veckok->v_dual.view_device().data();
888:     if (mtype) *mtype = PETSC_MEMTYPE_KOKKOS;
889:   }
890:   return 0;
891: }

893: PetscErrorCode VecRestoreArrayAndMemType_SeqKokkos(Vec v, PetscScalar **a)
894: {
895:   Vec_Kokkos *veckok = static_cast<Vec_Kokkos *>(v->spptr);

897:   if (std::is_same<DefaultMemorySpace, Kokkos::HostSpace>::value) {
898:     veckok->v_dual.modify_host();
899:   } else {
900:     veckok->v_dual.modify_device();
901:   }
902:   return 0;
903: }

905: PetscErrorCode VecGetArrayWriteAndMemType_SeqKokkos(Vec v, PetscScalar **a, PetscMemType *mtype)
906: {
907:   Vec_Kokkos *veckok = static_cast<Vec_Kokkos *>(v->spptr);

909:   if (std::is_same<DefaultMemorySpace, Kokkos::HostSpace>::value) {
910:     *a = veckok->v_dual.view_host().data();
911:     if (mtype) *mtype = PETSC_MEMTYPE_HOST;
912:   } else {
913:     /* When there is device, we always return device data (but no need to sync the device) */
914:     veckok->v_dual.clear_sync_state(); /* So that in restore, we can safely modify_device() */
915:     *a = veckok->v_dual.view_device().data();
916:     if (mtype) *mtype = PETSC_MEMTYPE_KOKKOS;
917:   }
918:   return 0;
919: }

921: /* Copy xin's sync state to y */
922: static PetscErrorCode VecCopySyncState_Kokkos_Private(Vec xin, Vec yout)
923: {
924:   Vec_Kokkos *xkok = static_cast<Vec_Kokkos *>(xin->spptr);
925:   Vec_Kokkos *ykok = static_cast<Vec_Kokkos *>(yout->spptr);

927:   ykok->v_dual.clear_sync_state();
928:   if (xkok->v_dual.need_sync_host()) {
929:     ykok->v_dual.modify_device();
930:   } else if (xkok->v_dual.need_sync_device()) {
931:     ykok->v_dual.modify_host();
932:   }
933:   return 0;
934: }

936: /* Interal routine shared by VecGetSubVector_{SeqKokkos,MPIKokkos} */
937: PetscErrorCode VecGetSubVector_Kokkos_Private(Vec x, PetscBool xIsMPI, IS is, Vec *y)
938: {
939:   PetscBool contig;
940:   PetscInt  n, N, start, bs;
941:   MPI_Comm  comm;
942:   Vec       z;

944:   PetscObjectGetComm((PetscObject)x, &comm);
945:   ISGetLocalSize(is, &n);
946:   ISGetSize(is, &N);
947:   VecGetSubVectorContiguityAndBS_Private(x, is, &contig, &start, &bs);

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

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

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

960:     /* This is relevant only in debug mode */
961:     PetscInt state = 0;
962:     VecLockGet(x, &state);
963:     if (state) { /* x is either in read or read/write mode, therefore z, overlapped with x, can only be in read mode */
964:       VecLockReadPush(z);
965:     }

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

970:   } else { /* Have to create a VecScatter and a stand-alone vector */
971:     VecGetSubVectorThroughVecScatter_Private(x, is, bs, &z);
972:   }
973:   *y = z;
974:   return 0;
975: }

977: PetscErrorCode VecGetSubVector_SeqKokkos(Vec x, IS is, Vec *y)
978: {
979:   VecGetSubVector_Kokkos_Private(x, PETSC_FALSE, is, y);
980:   return 0;
981: }

983: /* Restore subvector y to x */
984: PetscErrorCode VecRestoreSubVector_SeqKokkos(Vec x, IS is, Vec *y)
985: {
986:   VecScatter                    vscat;
987:   PETSC_UNUSED PetscObjectState dummystate = 0;
988:   PetscBool                     unchanged;

990:   PetscObjectComposedDataGetInt((PetscObject)*y, VecGetSubVectorSavedStateId, dummystate, unchanged);
991:   if (unchanged) return 0; /* If y's state has not changed since VecGetSubVector(), we only need to destroy it */

993:   PetscObjectQuery((PetscObject)*y, "VecGetSubVector_Scatter", (PetscObject *)&vscat);
994:   if (vscat) {
995:     VecScatterBegin(vscat, *y, x, INSERT_VALUES, SCATTER_REVERSE);
996:     VecScatterEnd(vscat, *y, x, INSERT_VALUES, SCATTER_REVERSE);
997:   } else { /* y and x's (host and device) arrays overlap */
998:     Vec_Kokkos *xkok = static_cast<Vec_Kokkos *>(x->spptr);
999:     Vec_Kokkos *ykok = static_cast<Vec_Kokkos *>((*y)->spptr);
1000:     PetscInt    state;

1002:     VecLockGet(x, &state);

1005:     /* The tricky part: one has to carefully sync the arrays */
1006:     if (xkok->v_dual.need_sync_device()) {      /* x's host has newer data */
1007:       ykok->v_dual.sync_host();                 /* Move y's latest values to host (since y is just a subset of x) */
1008:     } else if (xkok->v_dual.need_sync_host()) { /* x's device has newer data */
1009:       ykok->v_dual.sync_device();               /* Move y's latest data to device */
1010:     } else {                                    /* x's host and device data is already sync'ed; Copy y's sync state to x */
1011:       VecCopySyncState_Kokkos_Private(*y, x);
1012:     }
1013:     PetscObjectStateIncrease((PetscObject)x); /* Since x is updated */
1014:   }
1015:   return 0;
1016: }

1018: static PetscErrorCode VecSetPreallocationCOO_SeqKokkos(Vec x, PetscCount ncoo, const PetscInt coo_i[])
1019: {
1020:   Vec_Seq    *vecseq = static_cast<Vec_Seq *>(x->data);
1021:   Vec_Kokkos *veckok = static_cast<Vec_Kokkos *>(x->spptr);
1022:   PetscInt    m;

1024:   VecSetPreallocationCOO_Seq(x, ncoo, coo_i);
1025:   VecGetLocalSize(x, &m);
1026:   veckok->SetUpCOO(vecseq, m);
1027:   return 0;
1028: }

1030: static PetscErrorCode VecSetValuesCOO_SeqKokkos(Vec x, const PetscScalar v[], InsertMode imode)
1031: {
1032:   Vec_Seq                    *vecseq = static_cast<Vec_Seq *>(x->data);
1033:   Vec_Kokkos                 *veckok = static_cast<Vec_Kokkos *>(x->spptr);
1034:   const PetscCountKokkosView &jmap1  = veckok->jmap1_d;
1035:   const PetscCountKokkosView &perm1  = veckok->perm1_d;
1036:   PetscScalarKokkosView       xv; /* View for vector x */
1037:   ConstPetscScalarKokkosView  vv; /* View for array v[] */
1038:   PetscInt                    m;
1039:   PetscMemType                memtype;

1041:   VecGetLocalSize(x, &m);
1042:   PetscGetMemType(v, &memtype);
1043:   if (PetscMemTypeHost(memtype)) { /* If user gave v[] in host, we might need to copy it to device if any */
1044:     vv = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), ConstPetscScalarKokkosViewHost(v, vecseq->coo_n));
1045:   } else {
1046:     vv = ConstPetscScalarKokkosView(v, vecseq->coo_n); /* Directly use v[]'s memory */
1047:   }

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

1052:   Kokkos::parallel_for(
1053:     m, KOKKOS_LAMBDA(const PetscCount i) {
1054:       PetscScalar sum = 0.0;
1055:       for (PetscCount k = jmap1(i); k < jmap1(i + 1); k++) sum += vv(perm1(k));
1056:       xv(i) = (imode == INSERT_VALUES ? 0.0 : xv(i)) + sum;
1057:     });

1059:   if (imode == INSERT_VALUES) VecRestoreKokkosViewWrite(x, &xv);
1060:   else VecRestoreKokkosView(x, &xv);
1061:   return 0;
1062: }

1064: static PetscErrorCode VecSetOps_SeqKokkos(Vec v)
1065: {
1066:   v->ops->abs             = VecAbs_SeqKokkos;
1067:   v->ops->reciprocal      = VecReciprocal_SeqKokkos;
1068:   v->ops->pointwisemult   = VecPointwiseMult_SeqKokkos;
1069:   v->ops->min             = VecMin_SeqKokkos;
1070:   v->ops->max             = VecMax_SeqKokkos;
1071:   v->ops->sum             = VecSum_SeqKokkos;
1072:   v->ops->shift           = VecShift_SeqKokkos;
1073:   v->ops->norm            = VecNorm_SeqKokkos;
1074:   v->ops->scale           = VecScale_SeqKokkos;
1075:   v->ops->copy            = VecCopy_SeqKokkos;
1076:   v->ops->set             = VecSet_SeqKokkos;
1077:   v->ops->swap            = VecSwap_SeqKokkos;
1078:   v->ops->axpy            = VecAXPY_SeqKokkos;
1079:   v->ops->axpby           = VecAXPBY_SeqKokkos;
1080:   v->ops->axpbypcz        = VecAXPBYPCZ_SeqKokkos;
1081:   v->ops->pointwisedivide = VecPointwiseDivide_SeqKokkos;
1082:   v->ops->setrandom       = VecSetRandom_SeqKokkos;

1084:   v->ops->dot   = VecDot_SeqKokkos;
1085:   v->ops->tdot  = VecTDot_SeqKokkos;
1086:   v->ops->mdot  = VecMDot_SeqKokkos;
1087:   v->ops->mtdot = VecMTDot_SeqKokkos;

1089:   v->ops->dot_local   = VecDot_SeqKokkos;
1090:   v->ops->tdot_local  = VecTDot_SeqKokkos;
1091:   v->ops->mdot_local  = VecMDot_SeqKokkos;
1092:   v->ops->mtdot_local = VecMTDot_SeqKokkos;

1094:   v->ops->norm_local             = VecNorm_SeqKokkos;
1095:   v->ops->maxpy                  = VecMAXPY_SeqKokkos;
1096:   v->ops->aypx                   = VecAYPX_SeqKokkos;
1097:   v->ops->waxpy                  = VecWAXPY_SeqKokkos;
1098:   v->ops->dotnorm2               = VecDotNorm2_SeqKokkos;
1099:   v->ops->placearray             = VecPlaceArray_SeqKokkos;
1100:   v->ops->replacearray           = VecReplaceArray_SeqKokkos;
1101:   v->ops->resetarray             = VecResetArray_SeqKokkos;
1102:   v->ops->destroy                = VecDestroy_SeqKokkos;
1103:   v->ops->duplicate              = VecDuplicate_SeqKokkos;
1104:   v->ops->conjugate              = VecConjugate_SeqKokkos;
1105:   v->ops->getlocalvector         = VecGetLocalVector_SeqKokkos;
1106:   v->ops->restorelocalvector     = VecRestoreLocalVector_SeqKokkos;
1107:   v->ops->getlocalvectorread     = VecGetLocalVector_SeqKokkos;
1108:   v->ops->restorelocalvectorread = VecRestoreLocalVector_SeqKokkos;
1109:   v->ops->getarraywrite          = VecGetArrayWrite_SeqKokkos;
1110:   v->ops->getarray               = VecGetArray_SeqKokkos;
1111:   v->ops->restorearray           = VecRestoreArray_SeqKokkos;

1113:   v->ops->getarrayandmemtype      = VecGetArrayAndMemType_SeqKokkos;
1114:   v->ops->restorearrayandmemtype  = VecRestoreArrayAndMemType_SeqKokkos;
1115:   v->ops->getarraywriteandmemtype = VecGetArrayWriteAndMemType_SeqKokkos;
1116:   v->ops->getsubvector            = VecGetSubVector_SeqKokkos;
1117:   v->ops->restoresubvector        = VecRestoreSubVector_SeqKokkos;

1119:   v->ops->setpreallocationcoo = VecSetPreallocationCOO_SeqKokkos;
1120:   v->ops->setvaluescoo        = VecSetValuesCOO_SeqKokkos;
1121:   return 0;
1122: }

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

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

1130:   Level: beginner

1132: .seealso: `VecCreate()`, `VecSetType()`, `VecSetFromOptions()`, `VecCreateMPIWithArray()`, `VECMPI`, `VecType`, `VecCreateMPI()`
1133: M*/
1134: PetscErrorCode VecCreate_SeqKokkos(Vec v)
1135: {
1136:   Vec_Seq    *vecseq;
1137:   Vec_Kokkos *veckok;

1139:   PetscKokkosInitializeCheck();
1140:   PetscLayoutSetUp(v->map);
1141:   VecCreate_Seq(v); /* Build a sequential vector, allocate array */
1142:   PetscObjectChangeTypeName((PetscObject)v, VECSEQKOKKOS);
1143:   VecSetOps_SeqKokkos(v);

1146:   vecseq         = static_cast<Vec_Seq *>(v->data);
1147:   veckok         = new Vec_Kokkos(v->map->n, vecseq->array, NULL); /* Let host claim it has the latest data (zero) */
1148:   v->spptr       = static_cast<void *>(veckok);
1149:   v->offloadmask = PETSC_OFFLOAD_KOKKOS;
1150:   return 0;
1151: }

1153: /*@C
1154:    VecCreateSeqKokkosWithArray - Creates a Kokkos sequential array-style vector,
1155:    where the user provides the array space to store the vector values. The array
1156:    provided must be a device array.

1158:    Collective

1160:    Input Parameters:
1161: +  comm - the communicator, should be PETSC_COMM_SELF
1162: .  bs - the block size
1163: .  n - the vector length
1164: -  array - device memory where the vector elements are to be stored.

1166:    Output Parameter:
1167: .  v - the vector

1169:    Notes:
1170:    Use VecDuplicate() or VecDuplicateVecs() to form additional vectors of the
1171:    same type as an existing vector.

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

1176:    Level: intermediate

1178: .seealso: `VecCreateMPICUDAWithArray()`, `VecCreate()`, `VecDuplicate()`, `VecDuplicateVecs()`,
1179:           `VecCreateGhost()`, `VecCreateSeq()`, `VecCreateSeqWithArray()`,
1180:           `VecCreateMPIWithArray()`
1181: @*/
1182: PetscErrorCode VecCreateSeqKokkosWithArray(MPI_Comm comm, PetscInt bs, PetscInt n, const PetscScalar darray[], Vec *v)
1183: {
1184:   PetscMPIInt  size;
1185:   Vec          w;
1186:   Vec_Kokkos  *veckok = NULL;
1187:   PetscScalar *harray;

1189:   MPI_Comm_size(comm, &size);

1192:   PetscKokkosInitializeCheck();
1193:   VecCreate(comm, &w);
1194:   VecSetSizes(w, n, n);
1195:   VecSetBlockSize(w, bs);
1196:   if (!darray) { /* Allocate memory ourself if user provided NULL */
1197:     VecSetType(w, VECSEQKOKKOS);
1198:   } else {
1199:     /* Build a VECSEQ, get its harray, and then build Vec_Kokkos along with darray */
1200:     if (std::is_same<DefaultMemorySpace, Kokkos::HostSpace>::value) {
1201:       harray = const_cast<PetscScalar *>(darray);
1202:       VecCreate_Seq_Private(w, harray); /* Build a sequential vector with harray */
1203:     } else {
1204:       VecSetType(w, VECSEQ);
1205:       harray = static_cast<Vec_Seq *>(w->data)->array;
1206:     }
1207:     PetscObjectChangeTypeName((PetscObject)w, VECSEQKOKKOS); /* Change it to Kokkos */
1208:     VecSetOps_SeqKokkos(w);
1209:     veckok = new Vec_Kokkos(n, harray, const_cast<PetscScalar *>(darray));
1210:     veckok->v_dual.modify_device(); /* Mark the device is modified */
1211:     w->offloadmask = PETSC_OFFLOAD_KOKKOS;
1212:     w->spptr       = static_cast<void *>(veckok);
1213:   }
1214:   *v = w;
1215:   return 0;
1216: }

1218: /*
1219:    VecCreateSeqKokkosWithArrays_Private - Creates a Kokkos sequential array-style vector
1220:    with user-provided arrays on host and device.

1222:    Collective

1224:    Input Parameter:
1225: +  comm - the communicator, should be PETSC_COMM_SELF
1226: .  bs - the block size
1227: .  n - the vector length
1228: .  harray - host memory where the vector elements are to be stored.
1229: -  darray - device memory where the vector elements are to be stored.

1231:    Output Parameter:
1232: .  v - the vector

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

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

1242:    PETSc does NOT free the array when the vector is destroyed via VecDestroy().
1243:    Caller should not free the array until the vector is destroyed.
1244: */
1245: PetscErrorCode VecCreateSeqKokkosWithArrays_Private(MPI_Comm comm, PetscInt bs, PetscInt n, const PetscScalar harray[], const PetscScalar darray[], Vec *v)
1246: {
1247:   PetscMPIInt size;
1248:   Vec         w;

1250:   PetscKokkosInitializeCheck();
1251:   MPI_Comm_size(comm, &size);
1253:   if (n) {
1256:   }

1259:   VecCreateSeqWithArray(comm, bs, n, harray, &w);
1260:   PetscObjectChangeTypeName((PetscObject)w, VECSEQKOKKOS); /* Change it to Kokkos */
1261:   VecSetOps_SeqKokkos(w);
1262:   w->spptr = new Vec_Kokkos(n, const_cast<PetscScalar *>(harray), const_cast<PetscScalar *>(darray));
1263:   w->offloadmask = PETSC_OFFLOAD_KOKKOS;
1264:   *v             = w;
1265:   return 0;
1266: }

1268: /* TODO: ftn-auto generates veckok.kokkosf.c */
1269: /*@C
1270:  VecCreateSeqKokkos - Creates a standard, sequential array-style vector.

1272:  Collective

1274:  Input Parameter:
1275:  +  comm - the communicator, should be PETSC_COMM_SELF
1276:  -  n - the vector length

1278:  Output Parameter:
1279:  .  v - the vector

1281:  Notes:
1282:  Use VecDuplicate() or VecDuplicateVecs() to form additional vectors of the
1283:  same type as an existing vector.

1285:  Level: intermediate

1287:  .seealso: `VecCreateMPI()`, `VecCreate()`, `VecDuplicate()`, `VecDuplicateVecs()`, `VecCreateGhost()`
1288:  @*/
1289: PetscErrorCode VecCreateSeqKokkos(MPI_Comm comm, PetscInt n, Vec *v)
1290: {
1291:   PetscKokkosInitializeCheck();
1292:   VecCreate(comm, v);
1293:   VecSetSizes(*v, n, n);
1294:   VecSetType(*v, VECSEQKOKKOS); /* Calls VecCreate_SeqKokkos */
1295:   return 0;
1296: }

1298: /* Duplicate layout etc but not the values in the input vector */
1299: PetscErrorCode VecDuplicate_SeqKokkos(Vec win, Vec *v)
1300: {
1301:   VecDuplicate_Seq(win, v); /* It also dups ops of win */
1302:   return 0;
1303: }

1305: PetscErrorCode VecDestroy_SeqKokkos(Vec v)
1306: {
1307:   Vec_Kokkos *veckok = static_cast<Vec_Kokkos *>(v->spptr);
1308:   Vec_Seq    *vecseq = static_cast<Vec_Seq *>(v->data);

1310:   delete veckok;
1311:   v->spptr = NULL;
1312:   if (vecseq) VecDestroy_Seq(v);
1313:   return 0;
1314: }