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*/