Actual source code: dagetov.kokkos.cxx

  1: #include <petsc/private/vecimpl_kokkos.hpp>
  2: #include <petsc/private/dmdaimpl.h>
  3: #include <petscdmda_kokkos.hpp>

  5: /* SUBMANSEC = DMDA */

  7: /* Use macro instead of inlined function to avoid annoying warnings like: 'dof' may be used uninitialized in this function [-Wmaybe-uninitialized] */
  8: #define DMDA_VEC_GET_SHAPE(da, vec, xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof) \
  9:   do { \
 10:     PetscCall(DMDAGetCorners(da, &xs, &ys, &zs, &xm, &ym, &zm)); \
 11:     PetscCall(DMDAGetGhostCorners(da, &gxs, &gys, &gzs, &gxm, &gym, &gzm)); \
 12:     PetscCall(DMDAGetInfo(da, &dim, NULL, NULL, NULL, NULL, NULL, NULL, &dof, NULL, NULL, NULL, NULL, NULL)); \
 13:     /* Handle case where user passes in global vector as opposed to local */ \
 14:     PetscCall(VecGetLocalSize(vec, &N)); \
 15:     if (N == xm * ym * zm * dof) { \
 16:       gxm = xm; \
 17:       gym = ym; \
 18:       gzm = zm; \
 19:       gxs = xs; \
 20:       gys = ys; \
 21:       gzs = zs; \
 22:     } else PetscCheck(N == gxm * gym * gzm * dof, PETSC_COMM_SELF, PETSC_ERR_ARG_INCOMP, "Vector local size %" PetscInt_FMT " is not compatible with DMDA local sizes %" PetscInt_FMT " %" PetscInt_FMT, N, xm * ym * zm * dof, gxm * gym * gzm * dof); \
 23:   } while (0)

 25: /* -------------------- 1D ---------------- */
 26: template <class MemorySpace>
 27: PetscErrorCode DMDAVecGetKokkosOffsetView_Private(DM da, Vec vec, PetscScalarKokkosOffsetView1DType<MemorySpace> *ov, PetscBool overwrite)
 28: {
 29:   PetscInt                               xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof;
 30:   PetscScalarKokkosViewType<MemorySpace> kv;

 32:   PetscFunctionBegin;
 36:   DMDA_VEC_GET_SHAPE(da, vec, xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof);
 37:   PetscCheck(dim == 1, PETSC_COMM_SELF, PETSC_ERR_ARG_CORRUPT, "KokkosOffsetView is 1D but DMDA is %dD", (int)dim);
 38:   if (overwrite) PetscCall(VecGetKokkosViewWrite(vec, &kv));
 39:   else PetscCall(VecGetKokkosView(vec, &kv));
 40:   /* Construct the unmanaged OffsetView with {begin0,begin1,begins2},{end0,end1,end2} */
 41:   *ov = PetscScalarKokkosOffsetView1DType<MemorySpace>(kv.data(), {gxs * dof}, {(gxs + gxm) * dof});
 42:   PetscFunctionReturn(PETSC_SUCCESS);
 43: }

 45: template <class MemorySpace>
 46: PetscErrorCode DMDAVecRestoreKokkosOffsetView_Private(DM da, Vec vec, PetscScalarKokkosOffsetView1DType<MemorySpace> *ov, PetscBool overwrite)
 47: {
 48:   PetscScalarKokkosViewType<MemorySpace> kv;

 50:   PetscFunctionBegin;
 54:   kv = ov->view(); /* OffsetView to View */
 55:   if (overwrite) PetscCall(VecRestoreKokkosViewWrite(vec, &kv));
 56:   else PetscCall(VecRestoreKokkosView(vec, &kv));
 57:   PetscFunctionReturn(PETSC_SUCCESS);
 58: }

 60: template <class MemorySpace>
 61: PetscErrorCode DMDAVecGetKokkosOffsetView(DM da, Vec vec, ConstPetscScalarKokkosOffsetView1DType<MemorySpace> *ov)
 62: {
 63:   PetscInt                                    xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof;
 64:   ConstPetscScalarKokkosViewType<MemorySpace> kv;

 66:   PetscFunctionBegin;
 70:   DMDA_VEC_GET_SHAPE(da, vec, xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof);
 71:   PetscCheck(dim == 1, PETSC_COMM_SELF, PETSC_ERR_ARG_CORRUPT, "KokkosOffsetView is 1D but DMDA is %dD", (int)dim);
 72:   PetscCall(VecGetKokkosView(vec, &kv));
 73:   *ov = ConstPetscScalarKokkosOffsetView1DType<MemorySpace>(kv.data(), {gxs * dof}, {(gxs + gxm) * dof});
 74:   PetscFunctionReturn(PETSC_SUCCESS);
 75: }

 77: template <class MemorySpace>
 78: PetscErrorCode DMDAVecRestoreKokkosOffsetView(DM da, Vec vec, ConstPetscScalarKokkosOffsetView1DType<MemorySpace> *ov)
 79: {
 80:   ConstPetscScalarKokkosViewType<MemorySpace> kv;

 82:   PetscFunctionBegin;
 86:   kv = ov->view();
 87:   PetscCall(VecRestoreKokkosView(vec, &kv));
 88:   PetscFunctionReturn(PETSC_SUCCESS);
 89: }

 91: /* ============================== 2D ================================= */
 92: template <class MemorySpace>
 93: PetscErrorCode DMDAVecGetKokkosOffsetView_Private(DM da, Vec vec, PetscScalarKokkosOffsetView2DType<MemorySpace> *ov, PetscBool overwrite)
 94: {
 95:   PetscInt                               xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof;
 96:   PetscScalarKokkosViewType<MemorySpace> kv;

 98:   PetscFunctionBegin;
102:   DMDA_VEC_GET_SHAPE(da, vec, xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof);
103:   PetscCheck(dim == 2, PETSC_COMM_SELF, PETSC_ERR_ARG_CORRUPT, "KokkosOffsetView is 2D but DMDA is %dD", (int)dim);
104:   if (overwrite) PetscCall(VecGetKokkosViewWrite(vec, &kv));
105:   else PetscCall(VecGetKokkosView(vec, &kv));
106:   *ov = PetscScalarKokkosOffsetView2DType<MemorySpace>(kv.data(), {gys * dof, gxs * dof}, {(gys + gym) * dof, (gxs + gxm) * dof});
107:   PetscFunctionReturn(PETSC_SUCCESS);
108: }

110: template <class MemorySpace>
111: PetscErrorCode DMDAVecRestoreKokkosOffsetView_Private(DM da, Vec vec, PetscScalarKokkosOffsetView2DType<MemorySpace> *ov, PetscBool overwrite)
112: {
113:   PetscScalarKokkosViewType<MemorySpace> kv;

115:   PetscFunctionBegin;
119:   // kv   = ov->view(); /* 2D OffsetView => 2D View => 1D View. Why does it not work? */
120:   kv = PetscScalarKokkosViewType<MemorySpace>(ov->data(), ov->extent(0) * ov->extent(1));
121:   if (overwrite) PetscCall(VecRestoreKokkosViewWrite(vec, &kv));
122:   else PetscCall(VecRestoreKokkosView(vec, &kv));
123:   PetscFunctionReturn(PETSC_SUCCESS);
124: }

126: template <class MemorySpace>
127: PetscErrorCode DMDAVecGetKokkosOffsetView(DM da, Vec vec, ConstPetscScalarKokkosOffsetView2DType<MemorySpace> *ov)
128: {
129:   PetscInt                                    xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof;
130:   ConstPetscScalarKokkosViewType<MemorySpace> kv;

132:   PetscFunctionBegin;
136:   DMDA_VEC_GET_SHAPE(da, vec, xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof);
137:   PetscCheck(dim == 2, PETSC_COMM_SELF, PETSC_ERR_ARG_CORRUPT, "KokkosOffsetView is 2D but DMDA is %dD", (int)dim);
138:   PetscCall(VecGetKokkosView(vec, &kv));
139:   *ov = ConstPetscScalarKokkosOffsetView2DType<MemorySpace>(kv.data(), {gys * dof, gxs * dof}, {(gys + gym) * dof, (gxs + gxm) * dof});
140:   PetscFunctionReturn(PETSC_SUCCESS);
141: }

143: template <class MemorySpace>
144: PetscErrorCode DMDAVecRestoreKokkosOffsetView(DM da, Vec vec, ConstPetscScalarKokkosOffsetView2DType<MemorySpace> *ov)
145: {
146:   ConstPetscScalarKokkosViewType<MemorySpace> kv;

148:   PetscFunctionBegin;
152:   kv = ConstPetscScalarKokkosViewType<MemorySpace>(ov->data(), ov->extent(0) * ov->extent(1));
153:   PetscCall(VecRestoreKokkosView(vec, &kv));
154:   PetscFunctionReturn(PETSC_SUCCESS);
155: }

157: /* ============================== 3D ================================= */
158: template <class MemorySpace>
159: PetscErrorCode DMDAVecGetKokkosOffsetView_Private(DM da, Vec vec, PetscScalarKokkosOffsetView3DType<MemorySpace> *ov, PetscBool overwrite)
160: {
161:   PetscInt                               xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof;
162:   PetscScalarKokkosViewType<MemorySpace> kv;

164:   PetscFunctionBegin;
168:   DMDA_VEC_GET_SHAPE(da, vec, xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof);
169:   PetscCheck(dim == 3, PETSC_COMM_SELF, PETSC_ERR_ARG_CORRUPT, "KokkosOffsetView is 3D but DMDA is %dD", (int)dim);
170:   if (overwrite) PetscCall(VecGetKokkosViewWrite(vec, &kv));
171:   else PetscCall(VecGetKokkosView(vec, &kv));
172:   *ov = PetscScalarKokkosOffsetView3DType<MemorySpace>(kv.data(), {gzs * dof, gys * dof, gxs * dof}, {(gzs + gzm) * dof, (gys + gym) * dof, (gxs + gxm) * dof});
173:   PetscFunctionReturn(PETSC_SUCCESS);
174: }

176: template <class MemorySpace>
177: PetscErrorCode DMDAVecRestoreKokkosOffsetView_Private(DM da, Vec vec, PetscScalarKokkosOffsetView3DType<MemorySpace> *ov, PetscBool overwrite)
178: {
179:   PetscScalarKokkosViewType<MemorySpace> kv;

181:   PetscFunctionBegin;
185:   kv = PetscScalarKokkosViewType<MemorySpace>(ov->data(), ov->extent(0) * ov->extent(1) * ov->extent(2));
186:   if (overwrite) PetscCall(VecRestoreKokkosViewWrite(vec, &kv));
187:   else PetscCall(VecRestoreKokkosView(vec, &kv));
188:   PetscFunctionReturn(PETSC_SUCCESS);
189: }

191: template <class MemorySpace>
192: PetscErrorCode DMDAVecGetKokkosOffsetView(DM da, Vec vec, ConstPetscScalarKokkosOffsetView3DType<MemorySpace> *ov)
193: {
194:   PetscInt                                    xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof;
195:   ConstPetscScalarKokkosViewType<MemorySpace> kv;

197:   PetscFunctionBegin;
201:   DMDA_VEC_GET_SHAPE(da, vec, xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof);
202:   PetscCheck(dim == 3, PETSC_COMM_SELF, PETSC_ERR_ARG_CORRUPT, "KokkosOffsetView is 3D but DMDA is %dD", (int)dim);
203:   PetscCall(VecGetKokkosView(vec, &kv));
204:   *ov = ConstPetscScalarKokkosOffsetView3DType<MemorySpace>(kv.data(), {gzs * dof, gys * dof, gxs * dof}, {(gzs + gzm) * dof, (gys + gym) * dof, (gxs + gxm) * dof});
205:   PetscFunctionReturn(PETSC_SUCCESS);
206: }

208: template <class MemorySpace>
209: PetscErrorCode DMDAVecRestoreKokkosOffsetView(DM da, Vec vec, ConstPetscScalarKokkosOffsetView3DType<MemorySpace> *ov)
210: {
211:   ConstPetscScalarKokkosViewType<MemorySpace> kv;

213:   PetscFunctionBegin;
217:   kv = ConstPetscScalarKokkosViewType<MemorySpace>(ov->data(), ov->extent(0) * ov->extent(1) * ov->extent(2));
218:   PetscCall(VecRestoreKokkosView(vec, &kv));
219:   PetscFunctionReturn(PETSC_SUCCESS);
220: }

222: /* Function template explicit instantiation */
223: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetView(DM, Vec, ConstPetscScalarKokkosOffsetView1D *);
224: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetView(DM, Vec, ConstPetscScalarKokkosOffsetView1D *);
225: template <>
226: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetView(DM da, Vec vec, PetscScalarKokkosOffsetView1D *ov)
227: {
228:   return DMDAVecGetKokkosOffsetView_Private(da, vec, ov, PETSC_FALSE);
229: }
230: template <>
231: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetView(DM da, Vec vec, PetscScalarKokkosOffsetView1D *ov)
232: {
233:   return DMDAVecRestoreKokkosOffsetView_Private(da, vec, ov, PETSC_FALSE);
234: }
235: template <>
236: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewWrite(DM da, Vec vec, PetscScalarKokkosOffsetView1D *ov)
237: {
238:   return DMDAVecGetKokkosOffsetView_Private(da, vec, ov, PETSC_TRUE);
239: }
240: template <>
241: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewWrite(DM da, Vec vec, PetscScalarKokkosOffsetView1D *ov)
242: {
243:   return DMDAVecRestoreKokkosOffsetView_Private(da, vec, ov, PETSC_TRUE);
244: }

246: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetView(DM, Vec, ConstPetscScalarKokkosOffsetView2D *);
247: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetView(DM, Vec, ConstPetscScalarKokkosOffsetView2D *);
248: template <>
249: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetView(DM da, Vec vec, PetscScalarKokkosOffsetView2D *ov)
250: {
251:   return DMDAVecGetKokkosOffsetView_Private(da, vec, ov, PETSC_FALSE);
252: }
253: template <>
254: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetView(DM da, Vec vec, PetscScalarKokkosOffsetView2D *ov)
255: {
256:   return DMDAVecRestoreKokkosOffsetView_Private(da, vec, ov, PETSC_FALSE);
257: }
258: template <>
259: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewWrite(DM da, Vec vec, PetscScalarKokkosOffsetView2D *ov)
260: {
261:   return DMDAVecGetKokkosOffsetView_Private(da, vec, ov, PETSC_TRUE);
262: }
263: template <>
264: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewWrite(DM da, Vec vec, PetscScalarKokkosOffsetView2D *ov)
265: {
266:   return DMDAVecRestoreKokkosOffsetView_Private(da, vec, ov, PETSC_TRUE);
267: }

269: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetView(DM, Vec, ConstPetscScalarKokkosOffsetView3D *);
270: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetView(DM, Vec, ConstPetscScalarKokkosOffsetView3D *);
271: template <>
272: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetView(DM da, Vec vec, PetscScalarKokkosOffsetView3D *ov)
273: {
274:   return DMDAVecGetKokkosOffsetView_Private(da, vec, ov, PETSC_FALSE);
275: }
276: template <>
277: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetView(DM da, Vec vec, PetscScalarKokkosOffsetView3D *ov)
278: {
279:   return DMDAVecRestoreKokkosOffsetView_Private(da, vec, ov, PETSC_FALSE);
280: }
281: template <>
282: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewWrite(DM da, Vec vec, PetscScalarKokkosOffsetView3D *ov)
283: {
284:   return DMDAVecGetKokkosOffsetView_Private(da, vec, ov, PETSC_TRUE);
285: }
286: template <>
287: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewWrite(DM da, Vec vec, PetscScalarKokkosOffsetView3D *ov)
288: {
289:   return DMDAVecRestoreKokkosOffsetView_Private(da, vec, ov, PETSC_TRUE);
290: }

292: #if !defined(KOKKOS_ENABLE_DEFAULT_DEVICE_TYPE_HOST) /* Get host views if the default memory space is not host space */
293: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetView(DM, Vec, ConstPetscScalarKokkosOffsetView1DHost *);
294: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetView(DM, Vec, ConstPetscScalarKokkosOffsetView1DHost *);
295: template <>
296: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetView(DM da, Vec vec, PetscScalarKokkosOffsetView1DHost *ov)
297: {
298:   return DMDAVecGetKokkosOffsetView_Private(da, vec, ov, PETSC_FALSE);
299: }
300: template <>
301: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetView(DM da, Vec vec, PetscScalarKokkosOffsetView1DHost *ov)
302: {
303:   return DMDAVecRestoreKokkosOffsetView_Private(da, vec, ov, PETSC_FALSE);
304: }
305: template <>
306: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewWrite(DM da, Vec vec, PetscScalarKokkosOffsetView1DHost *ov)
307: {
308:   return DMDAVecGetKokkosOffsetView_Private(da, vec, ov, PETSC_TRUE);
309: }
310: template <>
311: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewWrite(DM da, Vec vec, PetscScalarKokkosOffsetView1DHost *ov)
312: {
313:   return DMDAVecRestoreKokkosOffsetView_Private(da, vec, ov, PETSC_TRUE);
314: }

316: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetView(DM, Vec, ConstPetscScalarKokkosOffsetView2DHost *);
317: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetView(DM, Vec, ConstPetscScalarKokkosOffsetView2DHost *);
318: template <>
319: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetView(DM da, Vec vec, PetscScalarKokkosOffsetView2DHost *ov)
320: {
321:   return DMDAVecGetKokkosOffsetView_Private(da, vec, ov, PETSC_FALSE);
322: }
323: template <>
324: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetView(DM da, Vec vec, PetscScalarKokkosOffsetView2DHost *ov)
325: {
326:   return DMDAVecRestoreKokkosOffsetView_Private(da, vec, ov, PETSC_FALSE);
327: }
328: template <>
329: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewWrite(DM da, Vec vec, PetscScalarKokkosOffsetView2DHost *ov)
330: {
331:   return DMDAVecGetKokkosOffsetView_Private(da, vec, ov, PETSC_TRUE);
332: }
333: template <>
334: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewWrite(DM da, Vec vec, PetscScalarKokkosOffsetView2DHost *ov)
335: {
336:   return DMDAVecRestoreKokkosOffsetView_Private(da, vec, ov, PETSC_TRUE);
337: }

339: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetView(DM, Vec, ConstPetscScalarKokkosOffsetView3DHost *);
340: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetView(DM, Vec, ConstPetscScalarKokkosOffsetView3DHost *);
341: template <>
342: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetView(DM da, Vec vec, PetscScalarKokkosOffsetView3DHost *ov)
343: {
344:   return DMDAVecGetKokkosOffsetView_Private(da, vec, ov, PETSC_FALSE);
345: }
346: template <>
347: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetView(DM da, Vec vec, PetscScalarKokkosOffsetView3DHost *ov)
348: {
349:   return DMDAVecRestoreKokkosOffsetView_Private(da, vec, ov, PETSC_FALSE);
350: }
351: template <>
352: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewWrite(DM da, Vec vec, PetscScalarKokkosOffsetView3DHost *ov)
353: {
354:   return DMDAVecGetKokkosOffsetView_Private(da, vec, ov, PETSC_TRUE);
355: }
356: template <>
357: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewWrite(DM da, Vec vec, PetscScalarKokkosOffsetView3DHost *ov)
358: {
359:   return DMDAVecRestoreKokkosOffsetView_Private(da, vec, ov, PETSC_TRUE);
360: }
361: #endif

363: /* ============================== 2D including DOF ================================= */
364: template <class MemorySpace>
365: PetscErrorCode DMDAVecGetKokkosOffsetViewDOF_Private(DM da, Vec vec, PetscScalarKokkosOffsetView2DType<MemorySpace> *ov, PetscBool overwrite)
366: {
367:   PetscInt                               xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof;
368:   PetscScalarKokkosViewType<MemorySpace> kv;

370:   PetscFunctionBegin;
374:   DMDA_VEC_GET_SHAPE(da, vec, xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof);
375:   PetscCheck(dim == 1, PETSC_COMM_SELF, PETSC_ERR_ARG_CORRUPT, "KokkosOffsetView is 2D but DMDA is %dD", (int)dim);
376:   if (overwrite) PetscCall(VecGetKokkosViewWrite(vec, &kv));
377:   else PetscCall(VecGetKokkosView(vec, &kv));
378:   *ov = PetscScalarKokkosOffsetView2DType<MemorySpace>(kv.data(), {gxs, 0}, {gxs + gxm, dof});
379:   PetscFunctionReturn(PETSC_SUCCESS);
380: }

382: template <class MemorySpace>
383: PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOF_Private(DM da, Vec vec, PetscScalarKokkosOffsetView2DType<MemorySpace> *ov, PetscBool overwrite)
384: {
385:   PetscScalarKokkosViewType<MemorySpace> kv;

387:   PetscFunctionBegin;
391:   kv = PetscScalarKokkosViewType<MemorySpace>(ov->data(), ov->extent(0) * ov->extent(1));
392:   if (overwrite) PetscCall(VecRestoreKokkosViewWrite(vec, &kv));
393:   else PetscCall(VecRestoreKokkosView(vec, &kv));
394:   PetscFunctionReturn(PETSC_SUCCESS);
395: }

397: template <class MemorySpace>
398: PetscErrorCode DMDAVecGetKokkosOffsetViewDOF(DM da, Vec vec, ConstPetscScalarKokkosOffsetView2DType<MemorySpace> *ov)
399: {
400:   PetscInt                                    xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof;
401:   ConstPetscScalarKokkosViewType<MemorySpace> kv;

403:   PetscFunctionBegin;
407:   DMDA_VEC_GET_SHAPE(da, vec, xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof);
408:   PetscCheck(dim == 1, PETSC_COMM_SELF, PETSC_ERR_ARG_CORRUPT, "KokkosOffsetView is 2D but DMDA is %dD", (int)dim);
409:   PetscCall(VecGetKokkosView(vec, &kv));
410:   *ov = ConstPetscScalarKokkosOffsetView2DType<MemorySpace>(kv.data(), {gxs, 0}, {gxs + gxm, dof});
411:   PetscFunctionReturn(PETSC_SUCCESS);
412: }

414: template <class MemorySpace>
415: PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOF(DM da, Vec vec, ConstPetscScalarKokkosOffsetView2DType<MemorySpace> *ov)
416: {
417:   ConstPetscScalarKokkosViewType<MemorySpace> kv;

419:   PetscFunctionBegin;
423:   kv = ConstPetscScalarKokkosViewType<MemorySpace>(ov->data(), ov->extent(0) * ov->extent(1));
424:   PetscCall(VecRestoreKokkosView(vec, &kv));
425:   PetscFunctionReturn(PETSC_SUCCESS);
426: }

428: /* ============================== 3D including DOF ================================= */
429: template <class MemorySpace>
430: PetscErrorCode DMDAVecGetKokkosOffsetViewDOF_Private(DM da, Vec vec, PetscScalarKokkosOffsetView3DType<MemorySpace> *ov, PetscBool overwrite)
431: {
432:   PetscInt                               xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof;
433:   PetscScalarKokkosViewType<MemorySpace> kv;

435:   PetscFunctionBegin;
439:   DMDA_VEC_GET_SHAPE(da, vec, xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof);
440:   PetscCheck(dim == 2, PETSC_COMM_SELF, PETSC_ERR_ARG_CORRUPT, "KokkosOffsetView is 3D but DMDA is %dD", (int)dim);
441:   if (overwrite) PetscCall(VecGetKokkosViewWrite(vec, &kv));
442:   else PetscCall(VecGetKokkosView(vec, &kv));
443:   *ov = PetscScalarKokkosOffsetView3DType<MemorySpace>(kv.data(), {gys, gxs, 0}, {gys + gym, gxs + gxm, dof});
444:   PetscFunctionReturn(PETSC_SUCCESS);
445: }

447: template <class MemorySpace>
448: PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOF_Private(DM da, Vec vec, PetscScalarKokkosOffsetView3DType<MemorySpace> *ov, PetscBool overwrite)
449: {
450:   PetscScalarKokkosViewType<MemorySpace> kv;

452:   PetscFunctionBegin;
456:   kv = PetscScalarKokkosViewType<MemorySpace>(ov->data(), ov->extent(0) * ov->extent(1) * ov->extent(2));
457:   if (overwrite) PetscCall(VecRestoreKokkosViewWrite(vec, &kv));
458:   else PetscCall(VecRestoreKokkosView(vec, &kv));
459:   PetscFunctionReturn(PETSC_SUCCESS);
460: }

462: template <class MemorySpace>
463: PetscErrorCode DMDAVecGetKokkosOffsetViewDOF(DM da, Vec vec, ConstPetscScalarKokkosOffsetView3DType<MemorySpace> *ov)
464: {
465:   PetscInt                                    xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof;
466:   ConstPetscScalarKokkosViewType<MemorySpace> kv;

468:   PetscFunctionBegin;
472:   DMDA_VEC_GET_SHAPE(da, vec, xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof);
473:   PetscCheck(dim == 2, PETSC_COMM_SELF, PETSC_ERR_ARG_CORRUPT, "KokkosOffsetView is 3D but DMDA is %dD", (int)dim);
474:   PetscCall(VecGetKokkosView(vec, &kv));
475:   *ov = ConstPetscScalarKokkosOffsetView3DType<MemorySpace>(kv.data(), {gys, gxs, 0}, {gys + gym, gxs + gxm, dof});
476:   PetscFunctionReturn(PETSC_SUCCESS);
477: }

479: template <class MemorySpace>
480: PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOF(DM da, Vec vec, ConstPetscScalarKokkosOffsetView3DType<MemorySpace> *ov)
481: {
482:   ConstPetscScalarKokkosViewType<MemorySpace> kv;

484:   PetscFunctionBegin;
488:   kv = ConstPetscScalarKokkosViewType<MemorySpace>(ov->data(), ov->extent(0) * ov->extent(1) * ov->extent(2));
489:   PetscCall(VecRestoreKokkosView(vec, &kv));
490:   PetscFunctionReturn(PETSC_SUCCESS);
491: }

493: /* ============================== 4D including DOF ================================= */
494: template <class MemorySpace>
495: PetscErrorCode DMDAVecGetKokkosOffsetViewDOF_Private(DM da, Vec vec, PetscScalarKokkosOffsetView4DType<MemorySpace> *ov, PetscBool overwrite)
496: {
497:   PetscInt                               xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof;
498:   PetscScalarKokkosViewType<MemorySpace> kv;

500:   PetscFunctionBegin;
504:   DMDA_VEC_GET_SHAPE(da, vec, xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof);
505:   PetscCheck(dim == 3, PETSC_COMM_SELF, PETSC_ERR_ARG_CORRUPT, "KokkosOffsetView is 4D but DMDA is %dD", (int)dim);
506:   if (overwrite) PetscCall(VecGetKokkosViewWrite(vec, &kv));
507:   else PetscCall(VecGetKokkosView(vec, &kv));
508:   *ov = PetscScalarKokkosOffsetView4DType<MemorySpace>(kv.data(), {gzs, gys, gxs, 0}, {gzs + gzm, gys + gym, gxs + gxm, dof});
509:   PetscFunctionReturn(PETSC_SUCCESS);
510: }

512: template <class MemorySpace>
513: PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOF_Private(DM da, Vec vec, PetscScalarKokkosOffsetView4DType<MemorySpace> *ov, PetscBool overwrite)
514: {
515:   PetscScalarKokkosViewType<MemorySpace> kv;

517:   PetscFunctionBegin;
521:   kv = PetscScalarKokkosViewType<MemorySpace>(ov->data(), ov->extent(0) * ov->extent(1) * ov->extent(2) * ov->extent(3));
522:   if (overwrite) PetscCall(VecRestoreKokkosViewWrite(vec, &kv));
523:   else PetscCall(VecRestoreKokkosView(vec, &kv));
524:   PetscFunctionReturn(PETSC_SUCCESS);
525: }

527: template <class MemorySpace>
528: PetscErrorCode DMDAVecGetKokkosOffsetViewDOF(DM da, Vec vec, ConstPetscScalarKokkosOffsetView4DType<MemorySpace> *ov)
529: {
530:   PetscInt                                    xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof;
531:   ConstPetscScalarKokkosViewType<MemorySpace> kv;

533:   PetscFunctionBegin;
537:   DMDA_VEC_GET_SHAPE(da, vec, xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof);
538:   PetscCheck(dim == 3, PETSC_COMM_SELF, PETSC_ERR_ARG_CORRUPT, "KokkosOffsetView is 4D but DMDA is %dD", (int)dim);
539:   PetscCall(VecGetKokkosView(vec, &kv));
540:   *ov = ConstPetscScalarKokkosOffsetView4DType<MemorySpace>(kv.data(), {gzs, gys, gxs, 0}, {gzs + gzm, gys + gym, gxs + gxm, dof});
541:   PetscFunctionReturn(PETSC_SUCCESS);
542: }

544: template <class MemorySpace>
545: PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOF(DM da, Vec vec, ConstPetscScalarKokkosOffsetView4DType<MemorySpace> *ov)
546: {
547:   ConstPetscScalarKokkosViewType<MemorySpace> kv;

549:   PetscFunctionBegin;
553:   kv = ConstPetscScalarKokkosViewType<MemorySpace>(ov->data(), ov->extent(0) * ov->extent(1) * ov->extent(2) * ov->extent(3));
554:   PetscCall(VecRestoreKokkosView(vec, &kv));
555:   PetscFunctionReturn(PETSC_SUCCESS);
556: }

558: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewDOF(DM, Vec, ConstPetscScalarKokkosOffsetView2D *);
559: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOF(DM, Vec, ConstPetscScalarKokkosOffsetView2D *);
560: template <>
561: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewDOF(DM da, Vec vec, PetscScalarKokkosOffsetView2D *ov)
562: {
563:   return DMDAVecGetKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_FALSE);
564: }
565: template <>
566: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOF(DM da, Vec vec, PetscScalarKokkosOffsetView2D *ov)
567: {
568:   return DMDAVecRestoreKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_FALSE);
569: }
570: template <>
571: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewDOFWrite(DM da, Vec vec, PetscScalarKokkosOffsetView2D *ov)
572: {
573:   return DMDAVecGetKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_TRUE);
574: }
575: template <>
576: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOFWrite(DM da, Vec vec, PetscScalarKokkosOffsetView2D *ov)
577: {
578:   return DMDAVecRestoreKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_TRUE);
579: }

581: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewDOF(DM, Vec, ConstPetscScalarKokkosOffsetView3D *);
582: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOF(DM, Vec, ConstPetscScalarKokkosOffsetView3D *);
583: template <>
584: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewDOF(DM da, Vec vec, PetscScalarKokkosOffsetView3D *ov)
585: {
586:   return DMDAVecGetKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_FALSE);
587: }
588: template <>
589: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOF(DM da, Vec vec, PetscScalarKokkosOffsetView3D *ov)
590: {
591:   return DMDAVecRestoreKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_FALSE);
592: }
593: template <>
594: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewDOFWrite(DM da, Vec vec, PetscScalarKokkosOffsetView3D *ov)
595: {
596:   return DMDAVecGetKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_TRUE);
597: }
598: template <>
599: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOFWrite(DM da, Vec vec, PetscScalarKokkosOffsetView3D *ov)
600: {
601:   return DMDAVecRestoreKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_TRUE);
602: }

604: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewDOF(DM, Vec, ConstPetscScalarKokkosOffsetView4D *);
605: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOF(DM, Vec, ConstPetscScalarKokkosOffsetView4D *);
606: template <>
607: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewDOF(DM da, Vec vec, PetscScalarKokkosOffsetView4D *ov)
608: {
609:   return DMDAVecGetKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_FALSE);
610: }
611: template <>
612: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOF(DM da, Vec vec, PetscScalarKokkosOffsetView4D *ov)
613: {
614:   return DMDAVecRestoreKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_FALSE);
615: }
616: template <>
617: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewDOFWrite(DM da, Vec vec, PetscScalarKokkosOffsetView4D *ov)
618: {
619:   return DMDAVecGetKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_TRUE);
620: }
621: template <>
622: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOFWrite(DM da, Vec vec, PetscScalarKokkosOffsetView4D *ov)
623: {
624:   return DMDAVecRestoreKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_TRUE);
625: }

627: #if !defined(KOKKOS_ENABLE_DEFAULT_DEVICE_TYPE_HOST) /* Get host views if the default memory space is not host space */
628: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewDOF(DM, Vec, ConstPetscScalarKokkosOffsetView2DHost *);
629: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOF(DM, Vec, ConstPetscScalarKokkosOffsetView2DHost *);
630: template <>
631: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewDOF(DM da, Vec vec, PetscScalarKokkosOffsetView2DHost *ov)
632: {
633:   return DMDAVecGetKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_FALSE);
634: }
635: template <>
636: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOF(DM da, Vec vec, PetscScalarKokkosOffsetView2DHost *ov)
637: {
638:   return DMDAVecRestoreKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_FALSE);
639: }
640: template <>
641: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewDOFWrite(DM da, Vec vec, PetscScalarKokkosOffsetView2DHost *ov)
642: {
643:   return DMDAVecGetKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_TRUE);
644: }
645: template <>
646: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOFWrite(DM da, Vec vec, PetscScalarKokkosOffsetView2DHost *ov)
647: {
648:   return DMDAVecRestoreKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_TRUE);
649: }

651: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewDOF(DM, Vec, ConstPetscScalarKokkosOffsetView3DHost *);
652: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOF(DM, Vec, ConstPetscScalarKokkosOffsetView3DHost *);
653: template <>
654: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewDOF(DM da, Vec vec, PetscScalarKokkosOffsetView3DHost *ov)
655: {
656:   return DMDAVecGetKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_FALSE);
657: }
658: template <>
659: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOF(DM da, Vec vec, PetscScalarKokkosOffsetView3DHost *ov)
660: {
661:   return DMDAVecRestoreKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_FALSE);
662: }
663: template <>
664: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewDOFWrite(DM da, Vec vec, PetscScalarKokkosOffsetView3DHost *ov)
665: {
666:   return DMDAVecGetKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_TRUE);
667: }
668: template <>
669: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOFWrite(DM da, Vec vec, PetscScalarKokkosOffsetView3DHost *ov)
670: {
671:   return DMDAVecRestoreKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_TRUE);
672: }

674: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewDOF(DM, Vec, ConstPetscScalarKokkosOffsetView4DHost *);
675: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOF(DM, Vec, ConstPetscScalarKokkosOffsetView4DHost *);
676: template <>
677: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewDOF(DM da, Vec vec, PetscScalarKokkosOffsetView4DHost *ov)
678: {
679:   return DMDAVecGetKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_FALSE);
680: }
681: template <>
682: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOF(DM da, Vec vec, PetscScalarKokkosOffsetView4DHost *ov)
683: {
684:   return DMDAVecRestoreKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_FALSE);
685: }
686: template <>
687: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewDOFWrite(DM da, Vec vec, PetscScalarKokkosOffsetView4DHost *ov)
688: {
689:   return DMDAVecGetKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_TRUE);
690: }
691: template <>
692: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOFWrite(DM da, Vec vec, PetscScalarKokkosOffsetView4DHost *ov)
693: {
694:   return DMDAVecRestoreKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_TRUE);
695: }
696: #endif