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