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;
35: PetscAssertPointer(ov, 3);
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;
53: PetscAssertPointer(ov, 3);
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;
69: PetscAssertPointer(ov, 3);
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;
85: PetscAssertPointer(ov, 3);
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;
101: PetscAssertPointer(ov, 3);
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;
118: PetscAssertPointer(ov, 3);
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;
135: PetscAssertPointer(ov, 3);
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;
151: PetscAssertPointer(ov, 3);
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;
167: PetscAssertPointer(ov, 3);
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;
184: PetscAssertPointer(ov, 3);
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;
200: PetscAssertPointer(ov, 3);
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;
216: PetscAssertPointer(ov, 3);
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;
373: PetscAssertPointer(ov, 3);
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;
390: PetscAssertPointer(ov, 3);
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;
406: PetscAssertPointer(ov, 3);
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;
422: PetscAssertPointer(ov, 3);
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;
438: PetscAssertPointer(ov, 3);
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;
455: PetscAssertPointer(ov, 3);
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;
471: PetscAssertPointer(ov, 3);
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;
487: PetscAssertPointer(ov, 3);
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;
503: PetscAssertPointer(ov, 3);
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;
520: PetscAssertPointer(ov, 3);
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;
536: PetscAssertPointer(ov, 3);
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;
552: PetscAssertPointer(ov, 3);
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