Actual source code: dagetov.kokkos.cxx
1: #include <petscdmda_kokkos.hpp>
2: #include <petsc/private/dmdaimpl.h>
3: #include <petsc/private/kokkosimpl.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 \
23: 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)); \
24: } while (0)
26: /* -------------------- 1D ---------------- */
27: template <class MemorySpace>
28: PetscErrorCode DMDAVecGetKokkosOffsetView_Private(DM da, Vec vec, PetscScalarKokkosOffsetView1DType<MemorySpace> *ov, PetscBool overwrite)
29: {
30: PetscInt xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof;
31: PetscScalarKokkosViewType<MemorySpace> kv;
33: PetscFunctionBegin;
36: PetscAssertPointer(ov, 3);
37: DMDA_VEC_GET_SHAPE(da, vec, xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof);
38: PetscCheck(dim == 1, PETSC_COMM_SELF, PETSC_ERR_ARG_CORRUPT, "KokkosOffsetView is 1D but DMDA is %dD", (int)dim);
39: if (overwrite) PetscCall(VecGetKokkosViewWrite(vec, &kv));
40: else PetscCall(VecGetKokkosView(vec, &kv));
41: /* Construct the unmanaged OffsetView with {begin0,begin1,begins2},{end0,end1,end2} */
42: *ov = PetscScalarKokkosOffsetView1DType<MemorySpace>(kv.data(), {gxs * dof}, {(gxs + gxm) * dof});
43: PetscFunctionReturn(PETSC_SUCCESS);
44: }
46: template <class MemorySpace>
47: PetscErrorCode DMDAVecRestoreKokkosOffsetView_Private(DM da, Vec vec, PetscScalarKokkosOffsetView1DType<MemorySpace> *ov, PetscBool overwrite)
48: {
49: PetscScalarKokkosViewType<MemorySpace> kv;
51: PetscFunctionBegin;
54: PetscAssertPointer(ov, 3);
55: kv = ov->view(); /* OffsetView to View */
56: if (overwrite) PetscCall(VecRestoreKokkosViewWrite(vec, &kv));
57: else PetscCall(VecRestoreKokkosView(vec, &kv));
58: PetscFunctionReturn(PETSC_SUCCESS);
59: }
61: template <class MemorySpace>
62: PetscErrorCode DMDAVecGetKokkosOffsetView(DM da, Vec vec, ConstPetscScalarKokkosOffsetView1DType<MemorySpace> *ov)
63: {
64: PetscInt xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof;
65: ConstPetscScalarKokkosViewType<MemorySpace> kv;
67: PetscFunctionBegin;
70: PetscAssertPointer(ov, 3);
71: DMDA_VEC_GET_SHAPE(da, vec, xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof);
72: PetscCheck(dim == 1, PETSC_COMM_SELF, PETSC_ERR_ARG_CORRUPT, "KokkosOffsetView is 1D but DMDA is %dD", (int)dim);
73: PetscCall(VecGetKokkosView(vec, &kv));
74: *ov = ConstPetscScalarKokkosOffsetView1DType<MemorySpace>(kv.data(), {gxs * dof}, {(gxs + gxm) * dof});
75: PetscFunctionReturn(PETSC_SUCCESS);
76: }
78: template <class MemorySpace>
79: PetscErrorCode DMDAVecRestoreKokkosOffsetView(DM da, Vec vec, ConstPetscScalarKokkosOffsetView1DType<MemorySpace> *ov)
80: {
81: ConstPetscScalarKokkosViewType<MemorySpace> kv;
83: PetscFunctionBegin;
86: PetscAssertPointer(ov, 3);
87: kv = ov->view();
88: PetscCall(VecRestoreKokkosView(vec, &kv));
89: PetscFunctionReturn(PETSC_SUCCESS);
90: }
92: /* ============================== 2D ================================= */
93: template <class MemorySpace>
94: PetscErrorCode DMDAVecGetKokkosOffsetView_Private(DM da, Vec vec, PetscScalarKokkosOffsetView2DType<MemorySpace> *ov, PetscBool overwrite)
95: {
96: PetscInt xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof;
97: PetscScalarKokkosViewType<MemorySpace> kv;
99: PetscFunctionBegin;
102: PetscAssertPointer(ov, 3);
103: DMDA_VEC_GET_SHAPE(da, vec, xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof);
104: PetscCheck(dim == 2, PETSC_COMM_SELF, PETSC_ERR_ARG_CORRUPT, "KokkosOffsetView is 2D but DMDA is %dD", (int)dim);
105: if (overwrite) PetscCall(VecGetKokkosViewWrite(vec, &kv));
106: else PetscCall(VecGetKokkosView(vec, &kv));
107: *ov = PetscScalarKokkosOffsetView2DType<MemorySpace>(kv.data(), {gys * dof, gxs * dof}, {(gys + gym) * dof, (gxs + gxm) * dof});
108: PetscFunctionReturn(PETSC_SUCCESS);
109: }
111: template <class MemorySpace>
112: PetscErrorCode DMDAVecRestoreKokkosOffsetView_Private(DM da, Vec vec, PetscScalarKokkosOffsetView2DType<MemorySpace> *ov, PetscBool overwrite)
113: {
114: PetscScalarKokkosViewType<MemorySpace> kv;
116: PetscFunctionBegin;
119: PetscAssertPointer(ov, 3);
120: // kv = ov->view(); /* 2D OffsetView => 2D View => 1D View. Why does it not work? */
121: kv = PetscScalarKokkosViewType<MemorySpace>(ov->data(), ov->extent(0) * ov->extent(1));
122: if (overwrite) PetscCall(VecRestoreKokkosViewWrite(vec, &kv));
123: else PetscCall(VecRestoreKokkosView(vec, &kv));
124: PetscFunctionReturn(PETSC_SUCCESS);
125: }
127: template <class MemorySpace>
128: PetscErrorCode DMDAVecGetKokkosOffsetView(DM da, Vec vec, ConstPetscScalarKokkosOffsetView2DType<MemorySpace> *ov)
129: {
130: PetscInt xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof;
131: ConstPetscScalarKokkosViewType<MemorySpace> kv;
133: PetscFunctionBegin;
136: PetscAssertPointer(ov, 3);
137: DMDA_VEC_GET_SHAPE(da, vec, xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof);
138: PetscCheck(dim == 2, PETSC_COMM_SELF, PETSC_ERR_ARG_CORRUPT, "KokkosOffsetView is 2D but DMDA is %dD", (int)dim);
139: PetscCall(VecGetKokkosView(vec, &kv));
140: *ov = ConstPetscScalarKokkosOffsetView2DType<MemorySpace>(kv.data(), {gys * dof, gxs * dof}, {(gys + gym) * dof, (gxs + gxm) * dof});
141: PetscFunctionReturn(PETSC_SUCCESS);
142: }
144: template <class MemorySpace>
145: PetscErrorCode DMDAVecRestoreKokkosOffsetView(DM da, Vec vec, ConstPetscScalarKokkosOffsetView2DType<MemorySpace> *ov)
146: {
147: ConstPetscScalarKokkosViewType<MemorySpace> kv;
149: PetscFunctionBegin;
152: PetscAssertPointer(ov, 3);
153: kv = ConstPetscScalarKokkosViewType<MemorySpace>(ov->data(), ov->extent(0) * ov->extent(1));
154: PetscCall(VecRestoreKokkosView(vec, &kv));
155: PetscFunctionReturn(PETSC_SUCCESS);
156: }
158: /* ============================== 3D ================================= */
159: template <class MemorySpace>
160: PetscErrorCode DMDAVecGetKokkosOffsetView_Private(DM da, Vec vec, PetscScalarKokkosOffsetView3DType<MemorySpace> *ov, PetscBool overwrite)
161: {
162: PetscInt xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof;
163: PetscScalarKokkosViewType<MemorySpace> kv;
165: PetscFunctionBegin;
168: PetscAssertPointer(ov, 3);
169: DMDA_VEC_GET_SHAPE(da, vec, xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof);
170: PetscCheck(dim == 3, PETSC_COMM_SELF, PETSC_ERR_ARG_CORRUPT, "KokkosOffsetView is 3D but DMDA is %dD", (int)dim);
171: if (overwrite) PetscCall(VecGetKokkosViewWrite(vec, &kv));
172: else PetscCall(VecGetKokkosView(vec, &kv));
173: *ov = PetscScalarKokkosOffsetView3DType<MemorySpace>(kv.data(), {gzs * dof, gys * dof, gxs * dof}, {(gzs + gzm) * dof, (gys + gym) * dof, (gxs + gxm) * dof});
174: PetscFunctionReturn(PETSC_SUCCESS);
175: }
177: template <class MemorySpace>
178: PetscErrorCode DMDAVecRestoreKokkosOffsetView_Private(DM da, Vec vec, PetscScalarKokkosOffsetView3DType<MemorySpace> *ov, PetscBool overwrite)
179: {
180: PetscScalarKokkosViewType<MemorySpace> kv;
182: PetscFunctionBegin;
185: PetscAssertPointer(ov, 3);
186: kv = PetscScalarKokkosViewType<MemorySpace>(ov->data(), ov->extent(0) * ov->extent(1) * ov->extent(2));
187: if (overwrite) PetscCall(VecRestoreKokkosViewWrite(vec, &kv));
188: else PetscCall(VecRestoreKokkosView(vec, &kv));
189: PetscFunctionReturn(PETSC_SUCCESS);
190: }
192: template <class MemorySpace>
193: PetscErrorCode DMDAVecGetKokkosOffsetView(DM da, Vec vec, ConstPetscScalarKokkosOffsetView3DType<MemorySpace> *ov)
194: {
195: PetscInt xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof;
196: ConstPetscScalarKokkosViewType<MemorySpace> kv;
198: PetscFunctionBegin;
201: PetscAssertPointer(ov, 3);
202: DMDA_VEC_GET_SHAPE(da, vec, xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof);
203: PetscCheck(dim == 3, PETSC_COMM_SELF, PETSC_ERR_ARG_CORRUPT, "KokkosOffsetView is 3D but DMDA is %dD", (int)dim);
204: PetscCall(VecGetKokkosView(vec, &kv));
205: *ov = ConstPetscScalarKokkosOffsetView3DType<MemorySpace>(kv.data(), {gzs * dof, gys * dof, gxs * dof}, {(gzs + gzm) * dof, (gys + gym) * dof, (gxs + gxm) * dof});
206: PetscFunctionReturn(PETSC_SUCCESS);
207: }
209: template <class MemorySpace>
210: PetscErrorCode DMDAVecRestoreKokkosOffsetView(DM da, Vec vec, ConstPetscScalarKokkosOffsetView3DType<MemorySpace> *ov)
211: {
212: ConstPetscScalarKokkosViewType<MemorySpace> kv;
214: PetscFunctionBegin;
217: PetscAssertPointer(ov, 3);
218: kv = ConstPetscScalarKokkosViewType<MemorySpace>(ov->data(), ov->extent(0) * ov->extent(1) * ov->extent(2));
219: PetscCall(VecRestoreKokkosView(vec, &kv));
220: PetscFunctionReturn(PETSC_SUCCESS);
221: }
223: /* Function template explicit instantiation */
224: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetView(DM, Vec, ConstPetscScalarKokkosOffsetView1D *);
225: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetView(DM, Vec, ConstPetscScalarKokkosOffsetView1D *);
226: template <>
227: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetView(DM da, Vec vec, PetscScalarKokkosOffsetView1D *ov)
228: {
229: return DMDAVecGetKokkosOffsetView_Private(da, vec, ov, PETSC_FALSE);
230: }
231: template <>
232: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetView(DM da, Vec vec, PetscScalarKokkosOffsetView1D *ov)
233: {
234: return DMDAVecRestoreKokkosOffsetView_Private(da, vec, ov, PETSC_FALSE);
235: }
236: template <>
237: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewWrite(DM da, Vec vec, PetscScalarKokkosOffsetView1D *ov)
238: {
239: return DMDAVecGetKokkosOffsetView_Private(da, vec, ov, PETSC_TRUE);
240: }
241: template <>
242: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewWrite(DM da, Vec vec, PetscScalarKokkosOffsetView1D *ov)
243: {
244: return DMDAVecRestoreKokkosOffsetView_Private(da, vec, ov, PETSC_TRUE);
245: }
247: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetView(DM, Vec, ConstPetscScalarKokkosOffsetView2D *);
248: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetView(DM, Vec, ConstPetscScalarKokkosOffsetView2D *);
249: template <>
250: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetView(DM da, Vec vec, PetscScalarKokkosOffsetView2D *ov)
251: {
252: return DMDAVecGetKokkosOffsetView_Private(da, vec, ov, PETSC_FALSE);
253: }
254: template <>
255: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetView(DM da, Vec vec, PetscScalarKokkosOffsetView2D *ov)
256: {
257: return DMDAVecRestoreKokkosOffsetView_Private(da, vec, ov, PETSC_FALSE);
258: }
259: template <>
260: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewWrite(DM da, Vec vec, PetscScalarKokkosOffsetView2D *ov)
261: {
262: return DMDAVecGetKokkosOffsetView_Private(da, vec, ov, PETSC_TRUE);
263: }
264: template <>
265: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewWrite(DM da, Vec vec, PetscScalarKokkosOffsetView2D *ov)
266: {
267: return DMDAVecRestoreKokkosOffsetView_Private(da, vec, ov, PETSC_TRUE);
268: }
270: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetView(DM, Vec, ConstPetscScalarKokkosOffsetView3D *);
271: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetView(DM, Vec, ConstPetscScalarKokkosOffsetView3D *);
272: template <>
273: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetView(DM da, Vec vec, PetscScalarKokkosOffsetView3D *ov)
274: {
275: return DMDAVecGetKokkosOffsetView_Private(da, vec, ov, PETSC_FALSE);
276: }
277: template <>
278: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetView(DM da, Vec vec, PetscScalarKokkosOffsetView3D *ov)
279: {
280: return DMDAVecRestoreKokkosOffsetView_Private(da, vec, ov, PETSC_FALSE);
281: }
282: template <>
283: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewWrite(DM da, Vec vec, PetscScalarKokkosOffsetView3D *ov)
284: {
285: return DMDAVecGetKokkosOffsetView_Private(da, vec, ov, PETSC_TRUE);
286: }
287: template <>
288: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewWrite(DM da, Vec vec, PetscScalarKokkosOffsetView3D *ov)
289: {
290: return DMDAVecRestoreKokkosOffsetView_Private(da, vec, ov, PETSC_TRUE);
291: }
293: #if !defined(KOKKOS_ENABLE_UNIFIED_MEMORY) /* Get host views if the default memory space is not host space */
294: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetView(DM, Vec, ConstPetscScalarKokkosOffsetView1DHost *);
295: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetView(DM, Vec, ConstPetscScalarKokkosOffsetView1DHost *);
296: template <>
297: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetView(DM da, Vec vec, PetscScalarKokkosOffsetView1DHost *ov)
298: {
299: return DMDAVecGetKokkosOffsetView_Private(da, vec, ov, PETSC_FALSE);
300: }
301: template <>
302: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetView(DM da, Vec vec, PetscScalarKokkosOffsetView1DHost *ov)
303: {
304: return DMDAVecRestoreKokkosOffsetView_Private(da, vec, ov, PETSC_FALSE);
305: }
306: template <>
307: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewWrite(DM da, Vec vec, PetscScalarKokkosOffsetView1DHost *ov)
308: {
309: return DMDAVecGetKokkosOffsetView_Private(da, vec, ov, PETSC_TRUE);
310: }
311: template <>
312: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewWrite(DM da, Vec vec, PetscScalarKokkosOffsetView1DHost *ov)
313: {
314: return DMDAVecRestoreKokkosOffsetView_Private(da, vec, ov, PETSC_TRUE);
315: }
317: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetView(DM, Vec, ConstPetscScalarKokkosOffsetView2DHost *);
318: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetView(DM, Vec, ConstPetscScalarKokkosOffsetView2DHost *);
319: template <>
320: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetView(DM da, Vec vec, PetscScalarKokkosOffsetView2DHost *ov)
321: {
322: return DMDAVecGetKokkosOffsetView_Private(da, vec, ov, PETSC_FALSE);
323: }
324: template <>
325: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetView(DM da, Vec vec, PetscScalarKokkosOffsetView2DHost *ov)
326: {
327: return DMDAVecRestoreKokkosOffsetView_Private(da, vec, ov, PETSC_FALSE);
328: }
329: template <>
330: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewWrite(DM da, Vec vec, PetscScalarKokkosOffsetView2DHost *ov)
331: {
332: return DMDAVecGetKokkosOffsetView_Private(da, vec, ov, PETSC_TRUE);
333: }
334: template <>
335: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewWrite(DM da, Vec vec, PetscScalarKokkosOffsetView2DHost *ov)
336: {
337: return DMDAVecRestoreKokkosOffsetView_Private(da, vec, ov, PETSC_TRUE);
338: }
340: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetView(DM, Vec, ConstPetscScalarKokkosOffsetView3DHost *);
341: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetView(DM, Vec, ConstPetscScalarKokkosOffsetView3DHost *);
342: template <>
343: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetView(DM da, Vec vec, PetscScalarKokkosOffsetView3DHost *ov)
344: {
345: return DMDAVecGetKokkosOffsetView_Private(da, vec, ov, PETSC_FALSE);
346: }
347: template <>
348: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetView(DM da, Vec vec, PetscScalarKokkosOffsetView3DHost *ov)
349: {
350: return DMDAVecRestoreKokkosOffsetView_Private(da, vec, ov, PETSC_FALSE);
351: }
352: template <>
353: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewWrite(DM da, Vec vec, PetscScalarKokkosOffsetView3DHost *ov)
354: {
355: return DMDAVecGetKokkosOffsetView_Private(da, vec, ov, PETSC_TRUE);
356: }
357: template <>
358: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewWrite(DM da, Vec vec, PetscScalarKokkosOffsetView3DHost *ov)
359: {
360: return DMDAVecRestoreKokkosOffsetView_Private(da, vec, ov, PETSC_TRUE);
361: }
362: #endif
364: /* ============================== 2D including DOF ================================= */
365: template <class MemorySpace>
366: PetscErrorCode DMDAVecGetKokkosOffsetViewDOF_Private(DM da, Vec vec, PetscScalarKokkosOffsetView2DType<MemorySpace> *ov, PetscBool overwrite)
367: {
368: PetscInt xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof;
369: PetscScalarKokkosViewType<MemorySpace> kv;
371: PetscFunctionBegin;
374: PetscAssertPointer(ov, 3);
375: DMDA_VEC_GET_SHAPE(da, vec, xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof);
376: PetscCheck(dim == 1, PETSC_COMM_SELF, PETSC_ERR_ARG_CORRUPT, "KokkosOffsetView is 2D but DMDA is %dD", (int)dim);
377: if (overwrite) PetscCall(VecGetKokkosViewWrite(vec, &kv));
378: else PetscCall(VecGetKokkosView(vec, &kv));
379: *ov = PetscScalarKokkosOffsetView2DType<MemorySpace>(kv.data(), {gxs, 0}, {gxs + gxm, dof});
380: PetscFunctionReturn(PETSC_SUCCESS);
381: }
383: template <class MemorySpace>
384: PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOF_Private(DM da, Vec vec, PetscScalarKokkosOffsetView2DType<MemorySpace> *ov, PetscBool overwrite)
385: {
386: PetscScalarKokkosViewType<MemorySpace> kv;
388: PetscFunctionBegin;
391: PetscAssertPointer(ov, 3);
392: kv = PetscScalarKokkosViewType<MemorySpace>(ov->data(), ov->extent(0) * ov->extent(1));
393: if (overwrite) PetscCall(VecRestoreKokkosViewWrite(vec, &kv));
394: else PetscCall(VecRestoreKokkosView(vec, &kv));
395: PetscFunctionReturn(PETSC_SUCCESS);
396: }
398: template <class MemorySpace>
399: PetscErrorCode DMDAVecGetKokkosOffsetViewDOF(DM da, Vec vec, ConstPetscScalarKokkosOffsetView2DType<MemorySpace> *ov)
400: {
401: PetscInt xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof;
402: ConstPetscScalarKokkosViewType<MemorySpace> kv;
404: PetscFunctionBegin;
407: PetscAssertPointer(ov, 3);
408: DMDA_VEC_GET_SHAPE(da, vec, xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof);
409: PetscCheck(dim == 1, PETSC_COMM_SELF, PETSC_ERR_ARG_CORRUPT, "KokkosOffsetView is 2D but DMDA is %dD", (int)dim);
410: PetscCall(VecGetKokkosView(vec, &kv));
411: *ov = ConstPetscScalarKokkosOffsetView2DType<MemorySpace>(kv.data(), {gxs, 0}, {gxs + gxm, dof});
412: PetscFunctionReturn(PETSC_SUCCESS);
413: }
415: template <class MemorySpace>
416: PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOF(DM da, Vec vec, ConstPetscScalarKokkosOffsetView2DType<MemorySpace> *ov)
417: {
418: ConstPetscScalarKokkosViewType<MemorySpace> kv;
420: PetscFunctionBegin;
423: PetscAssertPointer(ov, 3);
424: kv = ConstPetscScalarKokkosViewType<MemorySpace>(ov->data(), ov->extent(0) * ov->extent(1));
425: PetscCall(VecRestoreKokkosView(vec, &kv));
426: PetscFunctionReturn(PETSC_SUCCESS);
427: }
429: /* ============================== 3D including DOF ================================= */
430: template <class MemorySpace>
431: PetscErrorCode DMDAVecGetKokkosOffsetViewDOF_Private(DM da, Vec vec, PetscScalarKokkosOffsetView3DType<MemorySpace> *ov, PetscBool overwrite)
432: {
433: PetscInt xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof;
434: PetscScalarKokkosViewType<MemorySpace> kv;
436: PetscFunctionBegin;
439: PetscAssertPointer(ov, 3);
440: DMDA_VEC_GET_SHAPE(da, vec, xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof);
441: PetscCheck(dim == 2, PETSC_COMM_SELF, PETSC_ERR_ARG_CORRUPT, "KokkosOffsetView is 3D but DMDA is %dD", (int)dim);
442: if (overwrite) PetscCall(VecGetKokkosViewWrite(vec, &kv));
443: else PetscCall(VecGetKokkosView(vec, &kv));
444: *ov = PetscScalarKokkosOffsetView3DType<MemorySpace>(kv.data(), {gys, gxs, 0}, {gys + gym, gxs + gxm, dof});
445: PetscFunctionReturn(PETSC_SUCCESS);
446: }
448: template <class MemorySpace>
449: PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOF_Private(DM da, Vec vec, PetscScalarKokkosOffsetView3DType<MemorySpace> *ov, PetscBool overwrite)
450: {
451: PetscScalarKokkosViewType<MemorySpace> kv;
453: PetscFunctionBegin;
456: PetscAssertPointer(ov, 3);
457: kv = PetscScalarKokkosViewType<MemorySpace>(ov->data(), ov->extent(0) * ov->extent(1) * ov->extent(2));
458: if (overwrite) PetscCall(VecRestoreKokkosViewWrite(vec, &kv));
459: else PetscCall(VecRestoreKokkosView(vec, &kv));
460: PetscFunctionReturn(PETSC_SUCCESS);
461: }
463: template <class MemorySpace>
464: PetscErrorCode DMDAVecGetKokkosOffsetViewDOF(DM da, Vec vec, ConstPetscScalarKokkosOffsetView3DType<MemorySpace> *ov)
465: {
466: PetscInt xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof;
467: ConstPetscScalarKokkosViewType<MemorySpace> kv;
469: PetscFunctionBegin;
472: PetscAssertPointer(ov, 3);
473: DMDA_VEC_GET_SHAPE(da, vec, xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof);
474: PetscCheck(dim == 2, PETSC_COMM_SELF, PETSC_ERR_ARG_CORRUPT, "KokkosOffsetView is 3D but DMDA is %dD", (int)dim);
475: PetscCall(VecGetKokkosView(vec, &kv));
476: *ov = ConstPetscScalarKokkosOffsetView3DType<MemorySpace>(kv.data(), {gys, gxs, 0}, {gys + gym, gxs + gxm, dof});
477: PetscFunctionReturn(PETSC_SUCCESS);
478: }
480: template <class MemorySpace>
481: PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOF(DM da, Vec vec, ConstPetscScalarKokkosOffsetView3DType<MemorySpace> *ov)
482: {
483: ConstPetscScalarKokkosViewType<MemorySpace> kv;
485: PetscFunctionBegin;
488: PetscAssertPointer(ov, 3);
489: kv = ConstPetscScalarKokkosViewType<MemorySpace>(ov->data(), ov->extent(0) * ov->extent(1) * ov->extent(2));
490: PetscCall(VecRestoreKokkosView(vec, &kv));
491: PetscFunctionReturn(PETSC_SUCCESS);
492: }
494: /* ============================== 4D including DOF ================================= */
495: template <class MemorySpace>
496: PetscErrorCode DMDAVecGetKokkosOffsetViewDOF_Private(DM da, Vec vec, PetscScalarKokkosOffsetView4DType<MemorySpace> *ov, PetscBool overwrite)
497: {
498: PetscInt xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof;
499: PetscScalarKokkosViewType<MemorySpace> kv;
501: PetscFunctionBegin;
504: PetscAssertPointer(ov, 3);
505: DMDA_VEC_GET_SHAPE(da, vec, xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof);
506: PetscCheck(dim == 3, PETSC_COMM_SELF, PETSC_ERR_ARG_CORRUPT, "KokkosOffsetView is 4D but DMDA is %dD", (int)dim);
507: if (overwrite) PetscCall(VecGetKokkosViewWrite(vec, &kv));
508: else PetscCall(VecGetKokkosView(vec, &kv));
509: *ov = PetscScalarKokkosOffsetView4DType<MemorySpace>(kv.data(), {gzs, gys, gxs, 0}, {gzs + gzm, gys + gym, gxs + gxm, dof});
510: PetscFunctionReturn(PETSC_SUCCESS);
511: }
513: template <class MemorySpace>
514: PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOF_Private(DM da, Vec vec, PetscScalarKokkosOffsetView4DType<MemorySpace> *ov, PetscBool overwrite)
515: {
516: PetscScalarKokkosViewType<MemorySpace> kv;
518: PetscFunctionBegin;
521: PetscAssertPointer(ov, 3);
522: kv = PetscScalarKokkosViewType<MemorySpace>(ov->data(), ov->extent(0) * ov->extent(1) * ov->extent(2) * ov->extent(3));
523: if (overwrite) PetscCall(VecRestoreKokkosViewWrite(vec, &kv));
524: else PetscCall(VecRestoreKokkosView(vec, &kv));
525: PetscFunctionReturn(PETSC_SUCCESS);
526: }
528: template <class MemorySpace>
529: PetscErrorCode DMDAVecGetKokkosOffsetViewDOF(DM da, Vec vec, ConstPetscScalarKokkosOffsetView4DType<MemorySpace> *ov)
530: {
531: PetscInt xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof;
532: ConstPetscScalarKokkosViewType<MemorySpace> kv;
534: PetscFunctionBegin;
537: PetscAssertPointer(ov, 3);
538: DMDA_VEC_GET_SHAPE(da, vec, xs, ys, zs, xm, ym, zm, gxs, gys, gzs, gxm, gym, gzm, N, dim, dof);
539: PetscCheck(dim == 3, PETSC_COMM_SELF, PETSC_ERR_ARG_CORRUPT, "KokkosOffsetView is 4D but DMDA is %dD", (int)dim);
540: PetscCall(VecGetKokkosView(vec, &kv));
541: *ov = ConstPetscScalarKokkosOffsetView4DType<MemorySpace>(kv.data(), {gzs, gys, gxs, 0}, {gzs + gzm, gys + gym, gxs + gxm, dof});
542: PetscFunctionReturn(PETSC_SUCCESS);
543: }
545: template <class MemorySpace>
546: PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOF(DM da, Vec vec, ConstPetscScalarKokkosOffsetView4DType<MemorySpace> *ov)
547: {
548: ConstPetscScalarKokkosViewType<MemorySpace> kv;
550: PetscFunctionBegin;
553: PetscAssertPointer(ov, 3);
554: kv = ConstPetscScalarKokkosViewType<MemorySpace>(ov->data(), ov->extent(0) * ov->extent(1) * ov->extent(2) * ov->extent(3));
555: PetscCall(VecRestoreKokkosView(vec, &kv));
556: PetscFunctionReturn(PETSC_SUCCESS);
557: }
559: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewDOF(DM, Vec, ConstPetscScalarKokkosOffsetView2D *);
560: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOF(DM, Vec, ConstPetscScalarKokkosOffsetView2D *);
561: template <>
562: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewDOF(DM da, Vec vec, PetscScalarKokkosOffsetView2D *ov)
563: {
564: return DMDAVecGetKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_FALSE);
565: }
566: template <>
567: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOF(DM da, Vec vec, PetscScalarKokkosOffsetView2D *ov)
568: {
569: return DMDAVecRestoreKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_FALSE);
570: }
571: template <>
572: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewDOFWrite(DM da, Vec vec, PetscScalarKokkosOffsetView2D *ov)
573: {
574: return DMDAVecGetKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_TRUE);
575: }
576: template <>
577: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOFWrite(DM da, Vec vec, PetscScalarKokkosOffsetView2D *ov)
578: {
579: return DMDAVecRestoreKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_TRUE);
580: }
582: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewDOF(DM, Vec, ConstPetscScalarKokkosOffsetView3D *);
583: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOF(DM, Vec, ConstPetscScalarKokkosOffsetView3D *);
584: template <>
585: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewDOF(DM da, Vec vec, PetscScalarKokkosOffsetView3D *ov)
586: {
587: return DMDAVecGetKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_FALSE);
588: }
589: template <>
590: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOF(DM da, Vec vec, PetscScalarKokkosOffsetView3D *ov)
591: {
592: return DMDAVecRestoreKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_FALSE);
593: }
594: template <>
595: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewDOFWrite(DM da, Vec vec, PetscScalarKokkosOffsetView3D *ov)
596: {
597: return DMDAVecGetKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_TRUE);
598: }
599: template <>
600: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOFWrite(DM da, Vec vec, PetscScalarKokkosOffsetView3D *ov)
601: {
602: return DMDAVecRestoreKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_TRUE);
603: }
605: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewDOF(DM, Vec, ConstPetscScalarKokkosOffsetView4D *);
606: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOF(DM, Vec, ConstPetscScalarKokkosOffsetView4D *);
607: template <>
608: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewDOF(DM da, Vec vec, PetscScalarKokkosOffsetView4D *ov)
609: {
610: return DMDAVecGetKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_FALSE);
611: }
612: template <>
613: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOF(DM da, Vec vec, PetscScalarKokkosOffsetView4D *ov)
614: {
615: return DMDAVecRestoreKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_FALSE);
616: }
617: template <>
618: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewDOFWrite(DM da, Vec vec, PetscScalarKokkosOffsetView4D *ov)
619: {
620: return DMDAVecGetKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_TRUE);
621: }
622: template <>
623: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOFWrite(DM da, Vec vec, PetscScalarKokkosOffsetView4D *ov)
624: {
625: return DMDAVecRestoreKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_TRUE);
626: }
628: #if !defined(KOKKOS_ENABLE_UNIFIED_MEMORY) /* Get host views if the default memory space is not host space */
629: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewDOF(DM, Vec, ConstPetscScalarKokkosOffsetView2DHost *);
630: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOF(DM, Vec, ConstPetscScalarKokkosOffsetView2DHost *);
631: template <>
632: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewDOF(DM da, Vec vec, PetscScalarKokkosOffsetView2DHost *ov)
633: {
634: return DMDAVecGetKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_FALSE);
635: }
636: template <>
637: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOF(DM da, Vec vec, PetscScalarKokkosOffsetView2DHost *ov)
638: {
639: return DMDAVecRestoreKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_FALSE);
640: }
641: template <>
642: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewDOFWrite(DM da, Vec vec, PetscScalarKokkosOffsetView2DHost *ov)
643: {
644: return DMDAVecGetKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_TRUE);
645: }
646: template <>
647: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOFWrite(DM da, Vec vec, PetscScalarKokkosOffsetView2DHost *ov)
648: {
649: return DMDAVecRestoreKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_TRUE);
650: }
652: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewDOF(DM, Vec, ConstPetscScalarKokkosOffsetView3DHost *);
653: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOF(DM, Vec, ConstPetscScalarKokkosOffsetView3DHost *);
654: template <>
655: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewDOF(DM da, Vec vec, PetscScalarKokkosOffsetView3DHost *ov)
656: {
657: return DMDAVecGetKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_FALSE);
658: }
659: template <>
660: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOF(DM da, Vec vec, PetscScalarKokkosOffsetView3DHost *ov)
661: {
662: return DMDAVecRestoreKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_FALSE);
663: }
664: template <>
665: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewDOFWrite(DM da, Vec vec, PetscScalarKokkosOffsetView3DHost *ov)
666: {
667: return DMDAVecGetKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_TRUE);
668: }
669: template <>
670: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOFWrite(DM da, Vec vec, PetscScalarKokkosOffsetView3DHost *ov)
671: {
672: return DMDAVecRestoreKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_TRUE);
673: }
675: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewDOF(DM, Vec, ConstPetscScalarKokkosOffsetView4DHost *);
676: template PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOF(DM, Vec, ConstPetscScalarKokkosOffsetView4DHost *);
677: template <>
678: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewDOF(DM da, Vec vec, PetscScalarKokkosOffsetView4DHost *ov)
679: {
680: return DMDAVecGetKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_FALSE);
681: }
682: template <>
683: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOF(DM da, Vec vec, PetscScalarKokkosOffsetView4DHost *ov)
684: {
685: return DMDAVecRestoreKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_FALSE);
686: }
687: template <>
688: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecGetKokkosOffsetViewDOFWrite(DM da, Vec vec, PetscScalarKokkosOffsetView4DHost *ov)
689: {
690: return DMDAVecGetKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_TRUE);
691: }
692: template <>
693: PETSC_VISIBILITY_PUBLIC PetscErrorCode DMDAVecRestoreKokkosOffsetViewDOFWrite(DM da, Vec vec, PetscScalarKokkosOffsetView4DHost *ov)
694: {
695: return DMDAVecRestoreKokkosOffsetViewDOF_Private(da, vec, ov, PETSC_TRUE);
696: }
697: #endif