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:   PetscCall(MPIU_Allreduce(&lobj, obj, 1, MPIU_REAL, MPIU_SUM, comm));
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