Actual source code: ex55k.kokkos.cxx
1: #include <Kokkos_Core.hpp>
2: #include <petscdmda_kokkos.hpp>
4: #include <petscdm.h>
5: #include <petscdmda.h>
6: #include <petscsnes.h>
7: #include "ex55.h"
9: using DefaultMemorySpace = Kokkos::DefaultExecutionSpace::memory_space;
10: using ConstPetscScalarKokkosOffsetView2D = Kokkos::Experimental::OffsetView<const PetscScalar **, Kokkos::LayoutRight, DefaultMemorySpace>;
11: using PetscScalarKokkosOffsetView2D = Kokkos::Experimental::OffsetView<PetscScalar **, Kokkos::LayoutRight, DefaultMemorySpace>;
13: using PetscCountKokkosView = Kokkos::View<PetscCount *, DefaultMemorySpace>;
14: using PetscIntKokkosView = Kokkos::View<PetscInt *, DefaultMemorySpace>;
15: using PetscCountKokkosViewHost = Kokkos::View<PetscCount *, Kokkos::HostSpace>;
16: using PetscScalarKokkosView = Kokkos::View<PetscScalar *, DefaultMemorySpace>;
17: using Kokkos::Iterate;
18: using Kokkos::MDRangePolicy;
19: using Kokkos::Rank;
21: KOKKOS_INLINE_FUNCTION PetscErrorCode MMSSolution1(AppCtx *user, const DMDACoor2d *c, PetscScalar *u)
22: {
23: PetscReal x = PetscRealPart(c->x), y = PetscRealPart(c->y);
24: u[0] = x * (1 - x) * y * (1 - y);
25: return 0;
26: }
28: KOKKOS_INLINE_FUNCTION PetscErrorCode MMSForcing1(PetscReal user_param, const DMDACoor2d *c, PetscScalar *f)
29: {
30: PetscReal x = PetscRealPart(c->x), y = PetscRealPart(c->y);
31: f[0] = 2 * x * (1 - x) + 2 * y * (1 - y) - user_param * PetscExpReal(x * (1 - x) * y * (1 - y));
32: return 0;
33: }
35: PetscErrorCode FormFunctionLocalVec(DMDALocalInfo *info, Vec x, Vec f, AppCtx *user)
36: {
37: PetscReal lambda, hx, hy, hxdhy, hydhx;
38: PetscInt xs = info->xs, ys = info->ys, xm = info->xm, ym = info->ym, mx = info->mx, my = info->my;
39: PetscReal user_param = user->param;
41: ConstPetscScalarKokkosOffsetView2D xv;
42: PetscScalarKokkosOffsetView2D fv;
45: lambda = user->param;
46: hx = 1.0 / (PetscReal)(info->mx - 1);
47: hy = 1.0 / (PetscReal)(info->my - 1);
48: hxdhy = hx / hy;
49: hydhx = hy / hx;
50: /*
51: Compute function over the locally owned part of the grid
52: */
53: DMDAVecGetKokkosOffsetView(info->da, x, &xv);
54: DMDAVecGetKokkosOffsetViewWrite(info->da, f, &fv);
56: PetscCallCXX(Kokkos::parallel_for(
57: "FormFunctionLocalVec", MDRangePolicy<Rank<2, Iterate::Right, Iterate::Right>>({ys, xs}, {ys + ym, xs + xm}), KOKKOS_LAMBDA(PetscInt j, PetscInt i) {
58: DMDACoor2d c;
59: PetscScalar u, ue, uw, un, us, uxx, uyy, mms_solution, mms_forcing;
61: if (i == 0 || j == 0 || i == mx - 1 || j == my - 1) {
62: c.x = i * hx;
63: c.y = j * hy;
64: MMSSolution1(user, &c, &mms_solution);
65: fv(j, i) = 2.0 * (hydhx + hxdhy) * (xv(j, i) - mms_solution);
66: } else {
67: u = xv(j, i);
68: uw = xv(j, i - 1);
69: ue = xv(j, i + 1);
70: un = xv(j - 1, i);
71: us = xv(j + 1, i);
73: /* Enforce boundary conditions at neighboring points -- setting these values causes the Jacobian to be symmetric. */
74: if (i - 1 == 0) {
75: c.x = (i - 1) * hx;
76: c.y = j * hy;
77: MMSSolution1(user, &c, &uw);
78: }
79: if (i + 1 == mx - 1) {
80: c.x = (i + 1) * hx;
81: c.y = j * hy;
82: MMSSolution1(user, &c, &ue);
83: }
84: if (j - 1 == 0) {
85: c.x = i * hx;
86: c.y = (j - 1) * hy;
87: MMSSolution1(user, &c, &un);
88: }
89: if (j + 1 == my - 1) {
90: c.x = i * hx;
91: c.y = (j + 1) * hy;
92: MMSSolution1(user, &c, &us);
93: }
95: uxx = (2.0 * u - uw - ue) * hydhx;
96: uyy = (2.0 * u - un - us) * hxdhy;
97: mms_forcing = 0;
98: c.x = i * hx;
99: c.y = j * hy;
100: MMSForcing1(user_param, &c, &mms_forcing);
101: fv(j, i) = uxx + uyy - hx * hy * (lambda * PetscExpScalar(u) + mms_forcing);
102: }
103: }));
105: DMDAVecRestoreKokkosOffsetView(info->da, x, &xv);
106: DMDAVecRestoreKokkosOffsetViewWrite(info->da, f, &fv);
108: PetscLogFlops(11.0 * info->ym * info->xm);
109: return 0;
110: }
112: PetscErrorCode FormObjectiveLocalVec(DMDALocalInfo *info, Vec x, PetscReal *obj, AppCtx *user)
113: {
114: PetscInt xs = info->xs, ys = info->ys, xm = info->xm, ym = info->ym, mx = info->mx, my = info->my;
115: PetscReal lambda, hx, hy, hxdhy, hydhx, sc, lobj = 0;
116: MPI_Comm comm;
118: ConstPetscScalarKokkosOffsetView2D xv;
121: *obj = 0;
122: PetscObjectGetComm((PetscObject)info->da, &comm);
123: lambda = user->param;
124: hx = 1.0 / (PetscReal)(mx - 1);
125: hy = 1.0 / (PetscReal)(my - 1);
126: sc = hx * hy * lambda;
127: hxdhy = hx / hy;
128: hydhx = hy / hx;
129: /*
130: Compute function over the locally owned part of the grid
131: */
132: DMDAVecGetKokkosOffsetView(info->da, x, &xv);
134: PetscCallCXX(Kokkos::parallel_reduce(
135: "FormObjectiveLocalVec", MDRangePolicy<Rank<2, Iterate::Right, Iterate::Right>>({ys, xs}, {ys + ym, xs + xm}),
136: KOKKOS_LAMBDA(PetscInt j, PetscInt i, PetscReal & update) {
137: PetscScalar u, ue, uw, un, us, uxux, uyuy;
138: if (i == 0 || j == 0 || i == mx - 1 || j == my - 1) {
139: update += PetscRealPart((hydhx + hxdhy) * xv(j, i) * xv(j, i));
140: } else {
141: u = xv(j, i);
142: uw = xv(j, i - 1);
143: ue = xv(j, i + 1);
144: un = xv(j - 1, i);
145: us = xv(j + 1, i);
147: if (i - 1 == 0) uw = 0.;
148: if (i + 1 == mx - 1) ue = 0.;
149: if (j - 1 == 0) un = 0.;
150: if (j + 1 == my - 1) us = 0.;
152: /* F[u] = 1/2\int_{\omega}\nabla^2u(x)*u(x)*dx */
154: uxux = u * (2. * u - ue - uw) * hydhx;
155: uyuy = u * (2. * u - un - us) * hxdhy;
157: update += PetscRealPart(0.5 * (uxux + uyuy) - sc * PetscExpScalar(u));
158: }
159: },
160: lobj));
162: DMDAVecRestoreKokkosOffsetView(info->da, x, &xv);
163: PetscLogFlops(12.0 * info->ym * info->xm);
164: MPI_Allreduce(&lobj, obj, 1, MPIU_REAL, MPIU_SUM, comm);
165: return 0;
166: }
168: PetscErrorCode FormJacobianLocalVec(DMDALocalInfo *info, Vec x, Mat jac, Mat jacpre, AppCtx *user)
169: {
170: PetscInt i, j;
171: PetscInt xs = info->xs, ys = info->ys, xm = info->xm, ym = info->ym, mx = info->mx, my = info->my;
172: MatStencil col[5], row;
173: PetscScalar lambda, hx, hy, hxdhy, hydhx, sc;
174: DM coordDA;
175: Vec coordinates;
176: DMDACoor2d **coords;
179: lambda = user->param;
180: /* Extract coordinates */
181: DMGetCoordinateDM(info->da, &coordDA);
182: DMGetCoordinates(info->da, &coordinates);
184: DMDAVecGetArray(coordDA, coordinates, &coords);
185: hx = xm > 1 ? PetscRealPart(coords[ys][xs + 1].x) - PetscRealPart(coords[ys][xs].x) : 1.0;
186: hy = ym > 1 ? PetscRealPart(coords[ys + 1][xs].y) - PetscRealPart(coords[ys][xs].y) : 1.0;
187: DMDAVecRestoreArray(coordDA, coordinates, &coords);
189: hxdhy = hx / hy;
190: hydhx = hy / hx;
191: sc = hx * hy * lambda;
193: /* ----------------------------------------- */
194: /* MatSetPreallocationCOO() */
195: /* ----------------------------------------- */
196: PetscCount ncoo = ((PetscCount)xm) * ((PetscCount)ym) * 5;
197: PetscInt *coo_i, *coo_j, *ip, *jp;
198: PetscMalloc2(ncoo, &coo_i, ncoo, &coo_j); /* 5-point stencil such that each row has at most 5 nonzeros */
200: ip = coo_i;
201: jp = coo_j;
202: for (j = ys; j < ys + ym; j++) {
203: for (i = xs; i < xs + xm; i++) {
204: row.j = j;
205: row.i = i;
206: /* Initialize neighbors with negative indices */
207: col[0].j = col[1].j = col[3].j = col[4].j = -1;
208: /* boundary points */
209: if (i == 0 || j == 0 || i == mx - 1 || j == my - 1) {
210: col[2].j = row.j;
211: col[2].i = row.i;
212: } else {
213: /* interior grid points */
214: if (j - 1 != 0) {
215: col[0].j = j - 1;
216: col[0].i = i;
217: }
219: if (i - 1 != 0) {
220: col[1].j = j;
221: col[1].i = i - 1;
222: }
224: col[2].j = row.j;
225: col[2].i = row.i;
227: if (i + 1 != mx - 1) {
228: col[3].j = j;
229: col[3].i = i + 1;
230: }
232: if (j + 1 != mx - 1) {
233: col[4].j = j + 1;
234: col[4].i = i;
235: }
236: }
237: DMDAMapMatStencilToGlobal(info->da, 5, col, jp);
238: for (PetscInt k = 0; k < 5; k++) ip[k] = jp[2];
239: ip += 5;
240: jp += 5;
241: }
242: }
244: MatSetPreallocationCOO(jacpre, ncoo, coo_i, coo_j);
245: PetscFree2(coo_i, coo_j);
247: /* ----------------------------------------- */
248: /* MatSetValuesCOO() */
249: /* ----------------------------------------- */
250: PetscScalarKokkosView coo_v("coo_v", ncoo);
251: ConstPetscScalarKokkosOffsetView2D xv;
253: DMDAVecGetKokkosOffsetView(info->da, x, &xv);
255: PetscCallCXX(Kokkos::parallel_for(
256: "FormFunctionLocalVec", MDRangePolicy<Rank<2, Iterate::Right, Iterate::Right>>({ys, xs}, {ys + ym, xs + xm}), KOKKOS_LAMBDA(PetscCount j, PetscCount i) {
257: PetscInt p = ((j - ys) * xm + (i - xs)) * 5;
258: /* boundary points */
259: if (i == 0 || j == 0 || i == mx - 1 || j == my - 1) {
260: coo_v(p + 2) = 2.0 * (hydhx + hxdhy);
261: } else {
262: /* interior grid points */
263: if (j - 1 != 0) coo_v(p + 0) = -hxdhy;
264: if (i - 1 != 0) coo_v(p + 1) = -hydhx;
266: coo_v(p + 2) = 2.0 * (hydhx + hxdhy) - sc * PetscExpScalar(xv(j, i));
268: if (i + 1 != mx - 1) coo_v(p + 3) = -hydhx;
269: if (j + 1 != mx - 1) coo_v(p + 4) = -hxdhy;
270: }
271: }));
272: MatSetValuesCOO(jacpre, coo_v.data(), INSERT_VALUES);
273: DMDAVecRestoreKokkosOffsetView(info->da, x, &xv);
274: return 0;
275: }