Actual source code: ex55k.kokkos.cxx
1: #include <petscconf.h>
3: #if defined(PETSC_HAVE_KOKKOS_KERNELS)
4: #include <Kokkos_Core.hpp>
5: #include <petscdmda_kokkos.hpp>
7: #include <petscdm.h>
8: #include <petscdmda.h>
9: #include <petscsnes.h>
10: #include "ex55.h"
12: using DefaultMemorySpace = Kokkos::DefaultExecutionSpace::memory_space;
13: using ConstPetscScalarKokkosOffsetView2D = Kokkos::Experimental::OffsetView<const PetscScalar **, Kokkos::LayoutRight, DefaultMemorySpace>;
14: using PetscScalarKokkosOffsetView2D = Kokkos::Experimental::OffsetView<PetscScalar **, Kokkos::LayoutRight, DefaultMemorySpace>;
16: using PetscCountKokkosView = Kokkos::View<PetscCount *, DefaultMemorySpace>;
17: using PetscIntKokkosView = Kokkos::View<PetscInt *, DefaultMemorySpace>;
18: using PetscCountKokkosViewHost = Kokkos::View<PetscCount *, Kokkos::HostSpace>;
19: using PetscScalarKokkosView = Kokkos::View<PetscScalar *, DefaultMemorySpace>;
20: using Kokkos::Iterate;
21: using Kokkos::MDRangePolicy;
22: using Kokkos::Rank;
24: KOKKOS_INLINE_FUNCTION PetscErrorCode MMSSolution1(AppCtx *user, const DMDACoor2d *c, PetscScalar *u)
25: {
26: PetscReal x = PetscRealPart(c->x), y = PetscRealPart(c->y);
27: u[0] = x * (1 - x) * y * (1 - y);
28: return PETSC_SUCCESS;
29: }
31: KOKKOS_INLINE_FUNCTION PetscErrorCode MMSForcing1(PetscReal user_param, const DMDACoor2d *c, PetscScalar *f)
32: {
33: PetscReal x = PetscRealPart(c->x), y = PetscRealPart(c->y);
34: f[0] = 2 * x * (1 - x) + 2 * y * (1 - y) - user_param * PetscExpReal(x * (1 - x) * y * (1 - y));
35: return PETSC_SUCCESS;
36: }
38: PetscErrorCode FormFunctionLocalVec(DMDALocalInfo *info, Vec x, Vec f, AppCtx *user)
39: {
40: PetscReal lambda, hx, hy, hxdhy, hydhx;
41: PetscInt xs = info->xs, ys = info->ys, xm = info->xm, ym = info->ym, mx = info->mx, my = info->my;
42: PetscReal user_param = user->param;
44: ConstPetscScalarKokkosOffsetView2D xv;
45: PetscScalarKokkosOffsetView2D fv;
47: PetscFunctionBeginUser;
48: lambda = user->param;
49: hx = 1.0 / (PetscReal)(info->mx - 1);
50: hy = 1.0 / (PetscReal)(info->my - 1);
51: hxdhy = hx / hy;
52: hydhx = hy / hx;
53: /*
54: Compute function over the locally owned part of the grid
55: */
56: PetscCall(DMDAVecGetKokkosOffsetView(info->da, x, &xv));
57: PetscCall(DMDAVecGetKokkosOffsetViewWrite(info->da, f, &fv));
59: PetscCallCXX(Kokkos::parallel_for(
60: "FormFunctionLocalVec", MDRangePolicy<Rank<2, Iterate::Right, Iterate::Right>>({ys, xs}, {ys + ym, xs + xm}), KOKKOS_LAMBDA(PetscInt j, PetscInt i) {
61: DMDACoor2d c;
62: PetscScalar u, ue, uw, un, us, uxx, uyy, mms_solution, mms_forcing;
64: if (i == 0 || j == 0 || i == mx - 1 || j == my - 1) {
65: c.x = i * hx;
66: c.y = j * hy;
67: static_cast<void>(MMSSolution1(user, &c, &mms_solution));
68: fv(j, i) = 2.0 * (hydhx + hxdhy) * (xv(j, i) - mms_solution);
69: } else {
70: u = xv(j, i);
71: uw = xv(j, i - 1);
72: ue = xv(j, i + 1);
73: un = xv(j - 1, i);
74: us = xv(j + 1, i);
76: /* Enforce boundary conditions at neighboring points -- setting these values causes the Jacobian to be symmetric. */
77: if (i - 1 == 0) {
78: c.x = (i - 1) * hx;
79: c.y = j * hy;
80: static_cast<void>(MMSSolution1(user, &c, &uw));
81: }
82: if (i + 1 == mx - 1) {
83: c.x = (i + 1) * hx;
84: c.y = j * hy;
85: static_cast<void>(MMSSolution1(user, &c, &ue));
86: }
87: if (j - 1 == 0) {
88: c.x = i * hx;
89: c.y = (j - 1) * hy;
90: static_cast<void>(MMSSolution1(user, &c, &un));
91: }
92: if (j + 1 == my - 1) {
93: c.x = i * hx;
94: c.y = (j + 1) * hy;
95: static_cast<void>(MMSSolution1(user, &c, &us));
96: }
98: uxx = (2.0 * u - uw - ue) * hydhx;
99: uyy = (2.0 * u - un - us) * hxdhy;
100: mms_forcing = 0;
101: c.x = i * hx;
102: c.y = j * hy;
103: static_cast<void>(MMSForcing1(user_param, &c, &mms_forcing));
104: fv(j, i) = uxx + uyy - hx * hy * (lambda * PetscExpScalar(u) + mms_forcing);
105: }
106: }));
108: PetscCall(DMDAVecRestoreKokkosOffsetView(info->da, x, &xv));
109: PetscCall(DMDAVecRestoreKokkosOffsetViewWrite(info->da, f, &fv));
111: PetscCall(PetscLogFlops(11.0 * info->ym * info->xm));
112: PetscFunctionReturn(PETSC_SUCCESS);
113: }
115: PetscErrorCode FormObjectiveLocalVec(DMDALocalInfo *info, Vec x, PetscReal *obj, AppCtx *user)
116: {
117: PetscInt xs = info->xs, ys = info->ys, xm = info->xm, ym = info->ym, mx = info->mx, my = info->my;
118: PetscReal lambda, hx, hy, hxdhy, hydhx, sc, lobj = 0;
119: MPI_Comm comm;
121: ConstPetscScalarKokkosOffsetView2D xv;
123: PetscFunctionBeginUser;
124: *obj = 0;
125: PetscCall(PetscObjectGetComm((PetscObject)info->da, &comm));
126: lambda = user->param;
127: hx = 1.0 / (PetscReal)(mx - 1);
128: hy = 1.0 / (PetscReal)(my - 1);
129: sc = hx * hy * lambda;
130: hxdhy = hx / hy;
131: hydhx = hy / hx;
132: /*
133: Compute function over the locally owned part of the grid
134: */
135: PetscCall(DMDAVecGetKokkosOffsetView(info->da, x, &xv));
137: PetscCallCXX(Kokkos::parallel_reduce(
138: "FormObjectiveLocalVec", MDRangePolicy<Rank<2, Iterate::Right, Iterate::Right>>({ys, xs}, {ys + ym, xs + xm}),
139: KOKKOS_LAMBDA(PetscInt j, PetscInt i, PetscReal &update) {
140: PetscScalar u, ue, uw, un, us, uxux, uyuy;
141: if (i == 0 || j == 0 || i == mx - 1 || j == my - 1) {
142: update += PetscRealPart((hydhx + hxdhy) * xv(j, i) * xv(j, i));
143: } else {
144: u = xv(j, i);
145: uw = xv(j, i - 1);
146: ue = xv(j, i + 1);
147: un = xv(j - 1, i);
148: us = xv(j + 1, i);
150: if (i - 1 == 0) uw = 0.;
151: if (i + 1 == mx - 1) ue = 0.;
152: if (j - 1 == 0) un = 0.;
153: if (j + 1 == my - 1) us = 0.;
155: /* F[u] = 1/2\int_{\omega}\nabla^2u(x)*u(x)*dx */
157: uxux = u * (2. * u - ue - uw) * hydhx;
158: uyuy = u * (2. * u - un - us) * hxdhy;
160: update += PetscRealPart(0.5 * (uxux + uyuy) - sc * PetscExpScalar(u));
161: }
162: },
163: lobj));
165: PetscCall(DMDAVecRestoreKokkosOffsetView(info->da, x, &xv));
166: PetscCall(PetscLogFlops(12.0 * info->ym * info->xm));
167: *obj = lobj;
168: PetscFunctionReturn(PETSC_SUCCESS);
169: }
171: PetscErrorCode FormJacobianLocalVec(DMDALocalInfo *info, Vec x, Mat jac, Mat jacpre, AppCtx *user)
172: {
173: PetscInt i, j;
174: PetscInt xs = info->xs, ys = info->ys, xm = info->xm, ym = info->ym, mx = info->mx, my = info->my;
175: MatStencil col[5], row;
176: PetscScalar lambda, hx, hy, hxdhy, hydhx, sc;
177: DM coordDA;
178: Vec coordinates;
179: DMDACoor2d **coords;
181: PetscFunctionBeginUser;
182: lambda = user->param;
183: /* Extract coordinates */
184: PetscCall(DMGetCoordinateDM(info->da, &coordDA));
185: PetscCall(DMGetCoordinates(info->da, &coordinates));
187: PetscCall(DMDAVecGetArray(coordDA, coordinates, &coords));
188: hx = xm > 1 ? PetscRealPart(coords[ys][xs + 1].x) - PetscRealPart(coords[ys][xs].x) : 1.0;
189: hy = ym > 1 ? PetscRealPart(coords[ys + 1][xs].y) - PetscRealPart(coords[ys][xs].y) : 1.0;
190: PetscCall(DMDAVecRestoreArray(coordDA, coordinates, &coords));
192: hxdhy = hx / hy;
193: hydhx = hy / hx;
194: sc = hx * hy * lambda;
196: /* ----------------------------------------- */
197: /* MatSetPreallocationCOO() */
198: /* ----------------------------------------- */
199: if (!user->ncoo) {
200: PetscCount ncoo = ((PetscCount)xm) * ((PetscCount)ym) * 5;
201: PetscInt *coo_i, *coo_j, *ip, *jp;
202: PetscCall(PetscMalloc2(ncoo, &coo_i, ncoo, &coo_j)); /* 5-point stencil such that each row has at most 5 nonzeros */
204: ip = coo_i;
205: jp = coo_j;
206: for (j = ys; j < ys + ym; j++) {
207: for (i = xs; i < xs + xm; i++) {
208: row.j = j;
209: row.i = i;
210: /* Initialize neighbors with negative indices */
211: col[0].j = col[1].j = col[3].j = col[4].j = -1;
212: /* boundary points */
213: if (i == 0 || j == 0 || i == mx - 1 || j == my - 1) {
214: col[2].j = row.j;
215: col[2].i = row.i;
216: } else {
217: /* interior grid points */
218: if (j - 1 != 0) {
219: col[0].j = j - 1;
220: col[0].i = i;
221: }
223: if (i - 1 != 0) {
224: col[1].j = j;
225: col[1].i = i - 1;
226: }
228: col[2].j = row.j;
229: col[2].i = row.i;
231: if (i + 1 != mx - 1) {
232: col[3].j = j;
233: col[3].i = i + 1;
234: }
236: if (j + 1 != mx - 1) {
237: col[4].j = j + 1;
238: col[4].i = i;
239: }
240: }
241: PetscCall(DMDAMapMatStencilToGlobal(info->da, 5, col, jp));
242: for (PetscInt k = 0; k < 5; k++) ip[k] = jp[2];
243: ip += 5;
244: jp += 5;
245: }
246: }
247: PetscCall(MatSetPreallocationCOO(jacpre, ncoo, coo_i, coo_j));
248: PetscCall(PetscFree2(coo_i, coo_j));
249: user->ncoo = ncoo;
250: }
252: /* ----------------------------------------- */
253: /* MatSetValuesCOO() */
254: /* ----------------------------------------- */
255: PetscScalarKokkosView coo_v("coo_v", user->ncoo);
256: ConstPetscScalarKokkosOffsetView2D xv;
258: PetscCall(DMDAVecGetKokkosOffsetView(info->da, x, &xv));
260: PetscCallCXX(Kokkos::parallel_for(
261: "FormFunctionLocalVec", MDRangePolicy<Rank<2, Iterate::Right, Iterate::Right>>({ys, xs}, {ys + ym, xs + xm}), KOKKOS_LAMBDA(PetscCount j, PetscCount i) {
262: PetscInt p = ((j - ys) * xm + (i - xs)) * 5;
263: /* boundary points */
264: if (i == 0 || j == 0 || i == mx - 1 || j == my - 1) {
265: coo_v(p + 2) = 2.0 * (hydhx + hxdhy);
266: } else {
267: /* interior grid points */
268: if (j - 1 != 0) coo_v(p + 0) = -hxdhy;
269: if (i - 1 != 0) coo_v(p + 1) = -hydhx;
271: coo_v(p + 2) = 2.0 * (hydhx + hxdhy) - sc * PetscExpScalar(xv(j, i));
273: if (i + 1 != mx - 1) coo_v(p + 3) = -hydhx;
274: if (j + 1 != mx - 1) coo_v(p + 4) = -hxdhy;
275: }
276: }));
277: PetscCall(MatSetValuesCOO(jacpre, coo_v.data(), INSERT_VALUES));
278: PetscCall(DMDAVecRestoreKokkosOffsetView(info->da, x, &xv));
279: PetscFunctionReturn(PETSC_SUCCESS);
280: }
282: #else
283: #include "ex55.h"
285: PetscErrorCode FormObjectiveLocalVec(DMDALocalInfo *info, Vec x, PetscReal *obj, AppCtx *user)
286: {
287: PetscFunctionBeginUser;
288: PetscCheck(PETSC_FALSE, PETSC_COMM_SELF, PETSC_ERR_SUP, "Need to reconfigure with --download-kokkos-kernels");
289: PetscFunctionReturn(PETSC_SUCCESS);
290: }
292: PetscErrorCode FormFunctionLocalVec(DMDALocalInfo *info, Vec x, Vec f, AppCtx *user)
293: {
294: PetscFunctionBeginUser;
295: PetscCheck(PETSC_FALSE, PETSC_COMM_SELF, PETSC_ERR_SUP, "Need to reconfigure with --download-kokkos-kernels");
296: PetscFunctionReturn(PETSC_SUCCESS);
297: }
299: PetscErrorCode FormJacobianLocalVec(DMDALocalInfo *info, Vec x, Mat jac, Mat jacpre, AppCtx *user)
300: {
301: PetscFunctionBeginUser;
302: PetscCheck(PETSC_FALSE, PETSC_COMM_SELF, PETSC_ERR_SUP, "Need to reconfigure with --download-kokkos-kernels");
303: PetscFunctionReturn(PETSC_SUCCESS);
304: }
305: #endif
307: /*TEST
309: build:
310: TODO:
312: TEST*/